# save
torch.save(model.state_dict(), PATH)

# load
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

model.state_dict()其实返回的是一个OrderDict，存储了网络结构的名字和对应的参数。

在模型完成训练后，我们需要将训练好的模型保存为一个文件供测试使用，或者因为一些原因我们需要继续之前的状态训练之前保存的模型，那么如何在PyTorch中保存和恢复模型呢？

参考PyTorch官方的这份repo，我们知道有两种方法可以实现我们想要的效果。

方法一(推荐)：

第一种方法也是官方推荐的方法，只保存和恢复模型中的参数。

保存

恢复

使用这种方法，我们需要自己导入模型的结构信息。

备注：其中通过“torch.load（ XXX, map_location="cuda:0" if torch.cuda.is_available() else "cpu"） ” 将模型加载到GPU上 

方法二：

使用这种方法，将会保存模型的参数和结构信息。

保存

恢复

一个相对完整的例子

saving

torch.save({
            'epoch': epoch + 1,
            'arch': args.arch,
            'state_dict': model.state_dict(),
            'best_prec1': best_prec1,
        }, 'checkpoint.tar' )

loading

if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']
            model.load_state_dict(checkpoint['state_dict'])
            print("=> loaded checkpoint '{}' (epoch {})" .format(args.evaluate, checkpoint['epoch']))

获取模型中某些层的参数

对于恢复的模型，如果我们想查看某些层的参数，可以：

# 定义一个网络
from collections import OrderedDict
model = nn.Sequential(OrderedDict([
                  ('conv1', nn.Conv2d(1,20,5)),
                  ('relu1', nn.ReLU()),
                  ('conv2', nn.Conv2d(20,64,5)),
                  ('relu2', nn.ReLU())
                ]))
# 打印网络的结构
print(model)

Out:

Sequential (
  (conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
  (relu1): ReLU ()
  (conv2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
  (relu2): ReLU ()
)

如果我们想获取conv1的weight和bias：

params=model.state_dict() 
for k,v in params.items():
    print(k)    #打印网络中的变量名
print(params['conv1.weight'])   #打印conv1的weight
print(params['conv1.bias'])   #打印conv1的bias  

PyTorch 学习笔记（五）：存储和恢复模型并查看参数-PyTorch 中文网

源码详解Pytorch的state_dict和load_state_dict - 知乎

使用多张GPU进行训练以及推理调用模型

注意，使用GPU进行训练与推理的时，所用的代码不太一样：

训练是使用“ model.to(device) ”

推理时使用“ torch.load(r'./Gmodel.pth', map_location=device) ”

1、使用多张GPU进行训练的代码

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

model = getModel(args)

if torch.cuda.device_count() > 1:
	print("Let's use", torch.cuda.device_count(), "GPUs!")
	model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])   # device_ids=[0, 1, 2, 3]      
model.to(device)

注意：

1.当需要改变主GPU时，只需要修改如下：
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3"
device_ids=[0, 1, 2]  # 此时的GPU从GPU1开始，使用三张GPU
2.使用多张GPU时，batch_size大于1才能使得多张GPU都工作。

2、推理时

    G_model, _ = getModel(args)
    if torch.cuda.device_count() > 1:
        G_model = nn.DataParallel(G_model)
        checkpoint = torch.load(r'./Gmodel.pth', map_location=device)
        G_model.load_state_dict(checkpoint)
        if isinstance(G_model, torch.nn.DataParallel):
            G_model = G_model.module
    else:
        checkpoint = torch.load(r'./Gmodel.pth', map_location=device)
        G_model.load_state_dict(checkpoint)
    G_model = G_model.to(device)