Torch.load()使用方式

Posted 2023-04-14

参考技术A

torch.load() 的作用：从文件加载用 torch.save() 保存的对象。

api：

参数：

默认加载方式，使用cpu加载cpu训练得出的模型或者用gpu调用gpu训练的模型：

将全部 Tensor 全部加载到 cpu 上：

使用函数将所有张量加载到 CPU （适用在 GPU 训练的模型在 CPU 上加载）：

将所有张量加载到第一块 GPU (在 CPU 训练在 GPU 加载)：

将张量从 GPU 1 映射到 GPU 0 (第一块 GPU 训练，第二块 GPU 加载)：

根据你的设备，将张量加载到你当前设备上：

torch保存加载模型

• 保存模型

  torch.save(my_model.state_dict(), "params.pkl")

• 加载模型

  先初始化model网络结构
  model.load_state_dict(torch.load("params.pkl"))