PyTorch模型定义 | 模型容器 | 模型块 | 修改模型 | 模型读取与保存

一、PyTorch模型容器

  基于nn.Module，我们可以通过Sequential，ModuleList和ModuleDict三种方式定义PyTorch模型。

1.Sequential

   nn.Sequetial：按顺序的将一组网络层包装起来

其特性总结如下：

• 顺序性:各网络层之间严格按照顺序构建
• 自带forward():自带的forward里，通过for循环依次执行前向传播运算

下面使用Sequetial来包装LeNet模型。

class LeNetSequential(nn.Module):
    def __init__(self, classes):
        super(LeNetSequential, self).__init__()
        # 将卷积层与池化层包装成features网络层
        self.features = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),)
		
		# 将全连接网络包装成classifier层
        self.classifier = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, classes),)
            
	# 前向传播
    def forward(self, x):
        x = self.features(x)
        # 形状变换
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)
        return x

其源码为：

	class Sequential(Module):
	    r"""A sequential container.
	    Modules will be added to it in the order they are passed in the constructor.
	    Alternatively, an ordered dict of modules can also be passed in.
	
	    To make it easier to understand, here is a small example::
	
	        # Example of using Sequential
	        model = nn.Sequential(
	                  nn.Conv2d(1,20,5),
	                  nn.ReLU(),
	                  nn.Conv2d(20,64,5),
	                  nn.ReLU()
	                )
	
	        # Example of using Sequential with OrderedDict
	        model = nn.Sequential(OrderedDict([
	                  ('conv1', nn.Conv2d(1,20,5)),
	                  ('relu1', nn.ReLU()),
	                  ('conv2', nn.Conv2d(20,64,5)),
	                  ('relu2', nn.ReLU())
	                ]))
	    """
	
	    def __init__(self, *args):
	        super(Sequential, self).__init__()
	        # 将传入的网络层添加到Sequential中
	        # 判断输入的参数是否为有序字典
	        if len(args) == 1 and isinstance(args[0], OrderedDict):
	            for key, module in args[0].items():
	                self.add_module(key, module)
	        else:
	        # 如果不是的话，就直接将传入的网络层添加到Sequential中
	            for idx, module in enumerate(args):
	                self.add_module(str(idx), module)
	
	    def _get_item_by_idx(self, iterator, idx):
	        """Get the idx-th item of the iterator"""
	        size = len(self)
	        idx = operator.index(idx)
	        if not -size <= idx < size:
	            raise IndexError('index {} is out of range'.format(idx))
	        idx %= size
	        return next(islice(iterator, idx, None))
	
	    def __getitem__(self, idx):
	        if isinstance(idx, slice):
	            return self.__class__(OrderedDict(list(self._modules.items())[idx]))
	        else:
	            return self._get_item_by_idx(self._modules.values(), idx)
	
	    def __setitem__(self, idx, module):
	        key = self._get_item_by_idx(self._modules.keys(), idx)
	        return setattr(self, key, module)
	
	    def __delitem__(self, idx):
	        if isinstance(idx, slice):
	            for key in list(self._modules.keys())[idx]:
	                delattr(self, key)
	        else:
	            key = self._get_item_by_idx(self._modules.keys(), idx)
	            delattr(self, key)
	
	    def __len__(self):
	        return len(self._modules)
	
	    def __dir__(self):
	        keys = super(Sequential, self).__dir__()
	        keys = [key for key in keys if not key.isdigit()]
	        return keys
	
	    def forward(self, input):
	    	# 对Sequential中的网络层进行循环
	        for module in self._modules.values():
	            input = module(input)
	        return input

从源码中，我们可以发现Sequetial继承自Module，所以Sequetial仍然有8个有序字典。

LeNetSequential(
  (features): Sequential(
    (0): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
    (4): ReLU()
    (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (classifier): Sequential(
    (0): Linear(in_features=400, out_features=120, bias=True)
    (1): ReLU()
    (2): Linear(in_features=120, out_features=84, bias=True)
    (3): ReLU()
    (4): Linear(in_features=84, out_features=2, bias=True)
  )
)

打印上述网络层，我们可以发现上述网络层是没有命名的，是采用序号来索引网络层的。在深层网络中，很难采用序号来索引每个网络层。下面我们利用有序字典对网络层进行命名。

# 有序字典
from collections import OrderedDict

class LeNetSequentialOrderDict(nn.Module):
    def __init__(self, classes):
        super(LeNetSequentialOrderDict, self).__init__()

        self.features = nn.Sequential(OrderedDict({
            'conv1': nn.Conv2d(3, 6, 5),
            'relu1': nn.ReLU(inplace=True),
            'pool1': nn.MaxPool2d(kernel_size=2, stride=2),

            'conv2': nn.Conv2d(6, 16, 5),
            'relu2': nn.ReLU(inplace=True),
            'pool2': nn.MaxPool2d(kernel_size=2, stride=2),
        }))

        self.classifier = nn.Sequential(OrderedDict({
            'fc1': nn.Linear(16*5*5, 120),
            'relu3': nn.ReLU(),

            'fc2': nn.Linear(120, 84),
            'relu4': nn.ReLU(inplace=True),

            'fc3': nn.Linear(84, classes),
        }))

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)
        return x

此时，我们可以发现：网络层已经命名。我们可以通过名称来索引每个网络层。

	LeNetSequentialOrderDict(
	  (features): Sequential(
	    (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
	    (relu1): ReLU(inplace=True)
	    (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
	    (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
	    (relu2): ReLU(inplace=True)
	    (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
	  )
	  (classifier): Sequential(
	    (fc1): Linear(in_features=400, out_features=120, bias=True)
	    (relu3): ReLU()
	    (fc2): Linear(in_features=120, out_features=84, bias=True)
	    (relu4): ReLU(inplace=True)
	    (fc3): Linear(in_features=84, out_features=2, bias=True)
	  )
	)

可以看到，使用Sequential定义模型的好处在于简单、易读，同时使用Sequential定义的模型不需要再写forward，因为顺序已经定义好了。但使用Sequential也会使得模型定义丧失灵活性，比如需要在模型中间加入一个外部输入时就不适合用Sequential的方式实现。使用时需根据实际需求加以选择。

2.ModuleList

  ModuleList 接收一个子模块（或层，需属于nn.Module类）的列表作为输入，然后也可以类似List那样进行append和extend操作。同时，子模块或层的权重也会自动添加到网络中来。像python的list一样包装多个网络层，可以像python的list一样进行迭代，以迭代方式调用网络层。主要方法有：

• append():在ModuleList后面添加网络层
• extend():拼接两个ModuleList
• insert():指定在ModuleList中位置插入网络层

下面我们用nn.ModuleList来实现20层的全连接网络的实现

class ModuleList(nn.Module):
    def __init__(self):
        super(ModuleList, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(20)])

    def forward(self, x):
        for i, linear in enumerate(self.linears):
            x = linear(x)
        return x

要特别注意的是，nn.ModuleList 并没有定义一个网络，它只是将不同的模块储存在一起。ModuleList中元素的先后顺序并不代表其在网络中的真实位置顺序，需要经过forward函数指定各个层的先后顺序后才算完成了模型的定义。具体实现时用for循环即可完成。

3.ModuleDict

  ModuleDict和ModuleList的作用类似，只是ModuleDict能够更方便地为神经网络的层添加名称。像python的dict一样包装多个网络层，以索引方式调用网络层。主要方法有：

• clear():清空ModuleDict
• items():返回可迭代的键值对(key-value pairs)
• keys():返回字典的键(key)
• values():返回字典的值(value)
• pop():返回一对键值，并从字典中删除

class ModuleDict(nn.Module):
    def __init__(self):
        super(ModuleDict, self).__init__()
        self.choices = nn.ModuleDict({
            'conv': nn.Conv2d(10, 10, 3),
            'pool': nn.MaxPool2d(3)
        })

        self.activations = nn.ModuleDict({
            'relu': nn.ReLU(),
            'prelu': nn.PReLU()
        })
	
	# 两个可选择的属性
    def forward(self, x, choice, act):
        x = self.choices[choice](x)
        x = self.activations[act](x)
        return x

net = ModuleDict()
fake_img = torch.randn((4, 10, 32, 32))
output = net(fake_img, 'conv', 'relu')

4.三种容器的总结

• nn.Sequential :顺序性，各网络层之间严格按顺序执行，常用于block构建

• nn.ModuleList :迭代性，常用于大量重复网构建，通过for循环实现重复构建

• nn.ModuleDict :索引性,常用于可选择的网络层

• Sequential适用于快速验证结果，因为已经明确了要用哪些层，直接写一下就好了，不需要同时写__init__和forward；
• ModuleList和ModuleDict在某个完全相同的层需要重复出现多次时，非常方便实现，可以”一行顶多行“；
• 当我们需要之前层的信息的时候，比如 ResNets 中的残差计算，当前层的结果需要和之前层中的结果进行融合，一般使用 ModuleList/ModuleDict 比较方便。

5.案例：AlexNet构建

  AlexNet : 2012年以高出第二名10多个百分点的准确率获得lmageNet分类任务冠军，开创了卷积神经网络的新时代。AlexNet特点如下:

1. 采用ReLU：替换饱和激活函数，减轻梯度消失
2. 采用LRN(Local Response Normalization)：对数据归一化，减轻梯度消失
3. Dropout ：提高全连接层的鲁棒性,增加网络的泛化能力
4. Data Augmentation ：TenCrop，色彩修改

  AlexNet网络的结构如下：

  下面我们观察PyTorch提供的AlexNet网络的构建代码。

class AlexNet(nn.Module):

    def __init__(self, num_classes=1000):
        super(AlexNet, self).__init__()
        # features网络层用于特征提取
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        # 平均池化层
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        # classifier层用于分类
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )
	# 前向传播，因为使用了Sequential，使得非常简洁
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)

二、利用模型块快速搭建复杂网络

  对于大部分模型结构（比如ResNet、DenseNet等），我们仔细观察就会发现，虽然模型有很多层， 但是其中有很多重复出现的结构。考虑到每一层有其输入和输出，若干层串联成的”模块“也有其输入和输出，如果我们能将这些重复出现的层定义为一个”模块“，每次只需要向网络中添加对应的模块来构建模型，这样将会极大便利模型构建的过程。下面以U-Net为例，介绍如何构建模型块，以及如何利用模型块快速搭建复杂模型。

1.U-Net简介

  U-Net是分割 (Segmentation) 模型的杰作，在以医学影像为代表的诸多领域有着广泛的应用。U-Net模型结构如下图所示，通过残差连接结构解决了模型学习中的退化问题，使得神经网络的深度能够不断扩展。

2.U-Net模型块分析

  模型从上到下分为若干层，每层由左侧和右侧两个模型块组成，每侧的模型块与其上下模型块之间有连接；同时位于同一层左右两侧的模型块之间也有连接，称为“Skip-connection”。此外还有输入和输出处理等其他组成部分。由于模型的形状非常像英文字母的“U”，因此被命名为“U-Net”。
  组成U-Net的模型块主要有如下几个部分：

1. 每个子块内部的两次卷积（Double Convolution）
2. 左侧模型块之间的下采样连接，即最大池化（Max pooling）
3. 右侧模型块之间的上采样连接（Up sampling）
4. 输出层的处理

除模型块外，还有模型块之间的横向连接，输入和U-Net底部的连接等计算，这些单独的操作可以通过forward函数来实现。

3.U-Net模型块实现

  下面我们用PyTorch先实现上述的模型块，然后再利用定义好的模型块构建U-Net模型。先定义好各模型块，再定义模型块之间的连接顺序和计算方式。就好比装配零件一样，我们先装配好一些基础的部件，之后再用这些可以复用的部件得到整个装配体。这里的基础部件对应的四个模型块，根据功能我们将其命名为：DoubleConv，Down， Up， OutConv。

import torch
import torch.nn as nn
import torch.nn.functional as F

每个子块内部的两次卷积：

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

下采样连接：

class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)

上采样连接：

class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=False):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])
        # if you have padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

4.利用模型块组装U-Net

  使用写好的模型块，可以非常方便地组装U-Net模型。可以看到，通过模型块的方式实现了代码复用，整个模型结构定义所需的代码总行数明显减少，代码可读性也得到了提升。

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=False):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        factor = 2 if bilinear else 1
        self.down4 = Down(512, 1024 // factor)
        self.up1 = Up(1024, 512 // factor, bilinear)
        self.up2 = Up(512, 256 // factor, bilinear)
        self.up3 = Up(256, 128 // factor, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

三、修改模型

  除了自己构建PyTorch模型外，还有另一种应用场景：我们已经有一个现成的模型，但该模型中的部分结构不符合我们的要求，为了使用模型，我们需要对模型结构进行必要的修改。随着深度学习的发展和PyTorch越来越广泛的使用，有越来越多的开源模型可以供我们使用，很多时候我们也不必从头开始构建模型。因此，掌握如何修改PyTorch模型就显得尤为重要。

1.修改模型若干层

  我们这里以pytorch官方视觉库torchvision预定义好的模型ResNet50为例，探索如何修改模型的某一层或者某几层。我们先看看模型的定义是怎样的：

import torchvision.models as models
net = models.resnet50()
print(net)

ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
    (0): Bottleneck(
      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (downsample): Sequential(
        (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
..............
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=2048, out_features=1000, bias=True)
)

这里模型结构是为了适配ImageNet预训练的权重，因此最后全连接层fc的输出节点数是1000。假设我们要用这个resnet模型去做一个10分类的问题，就应该修改模型的fc层，将其输出节点数替换为10。另外，我们觉得一层全连接层可能太少了，想再加一层。可以做如下修改：

from collections import OrderedDict
classifier = nn.Sequential(OrderedDict([('fc1', nn.Linear(2048, 128)),
                          ('relu1', nn.ReLU()), 
                          ('dropout1',nn.Dropout(0.5)),
                          ('fc2', nn.Linear(128, 10)),
                          ('output', nn.Softmax(dim=1))
                          ]))
    
net.fc = classifier

这里的操作相当于将模型net最后名称为“fc”的层替换成了名称为“classifier”的结构，该结构是我们自己定义的。这里使用Sequential+OrderedDict的模型定义方式。至此，我们就完成了模型的修改，现在的模型就可以去做10分类任务了。

2.添加外部输入

  有时候在模型训练中，除了已有模型的输入之外，还需要输入额外的信息。比如在CNN网络中，我们除了输入图像，还需要同时输入图像对应的其他信息，这时候就需要在已有的CNN网络中添加额外的输入变量。基本思路是：将原模型添加输入位置前的部分作为一个整体，同时在forward中定义好原模型不变的部分、添加的输入和后续层之间的连接关系，从而完成模型的修改。
  我们以torchvision的resnet50模型为基础，任务还是10分类任务。不同点在于，我们希望利用已有的模型结构，在倒数第二层增加一个额外的输入变量add_variable来辅助预测。具体实现如下：

class Model(nn.Module):
    def __init__(self, net):
        super(Model, self).__init__()
        self.net = net
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc_add = nn.Linear(1001, 10, bias=True)
        self.output = nn.Softmax(dim=1)
        
    def forward(self, x, add_variable):
        x = self.net(x)
        x = torch.cat((self.dropout(self.relu(x)), add_variable.unsqueeze(1)),1)
        x = self.fc_add(x)
        x = self.output(x)
        return x

  这里的实现要点是通过torch.cat实现了tensor的拼接。torchvision中的resnet50输出是一个1000维的tensor，我们通过修改forward函数（配套定义一些层），先将1000维的tensor通过激活函数层和dropout层，再和外部输入变量"add_variable"拼接，最后通过全连接层映射到指定的输出维度10。
  另外这里对外部输入变量"add_variable"进行unsqueeze操作是为了和net输出的tensor保持维度一致，常用于add_variable是单一数值 (scalar) 的情况，此时add_variable的维度是 (batch_size, )，需要在第二维补充维数1，从而可以和tensor进行torch.cat操作。
  之后对我们修改好的模型结构进行实例化，就可以使用了：

import torchvision.models as models
net = models.resnet50()
model = Model(net).cuda()

  另外别忘了，训练中在输入数据的时候要给两个inputs：

outputs = model(inputs, add_var)

3.添加额外输出

  有时候在模型训练中，除了模型最后的输出外，我们需要输出模型某一中间层的结果，以施加额外的监督，获得更好的中间层结果。基本思路是修改模型定义中forward函数的return变量。
  我们依然以resnet50做10分类任务为例，在已经定义好的模型结构上，同时输出1000维的倒数第二层和10维的最后一层结果。具体实现如下：

class Model(nn.Module):
    def __init__(self, net):
        super(Model, self).__init__()
        self.net = net
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc1 = nn.Linear(1000, 10, bias=True)
        self.output = nn.Softmax(dim=1)
        
    def forward(self, x, add_variable):
        x1000 = self.net(x)
        x10 = self.dropout(self.relu(x1000))
        x10 = self.fc1(x10)
        x10 = self.output(x10)
        return x10, x1000

  对我们修改好的模型结构进行实例化，就可以使用了：

import torchvision.models as models
net = models.resnet50()
model = Model(net).cuda()

另外别忘了，训练中在输入数据后会有两个outputs：

out10, out1000 = model(inputs, add_var)

四、PyTorch的模型读取与保存

1.序列化与反序列化

  训练好的模型存储在内存中，内存中的数据不具备长久性存储的功能，所以，我们需要将模型从内存搬到硬盘中进行长久的存储，以备将来使用。这就是模型的保存与加载，也即序列化与反序列化。序列化指的是将内存中的对象保存到硬盘当中，以二进制序列的形式存储下来。反序列化指的是将硬盘当中的二进制序列反序列化的存储到内存中，得到对象，这样，我们就可以在内存中使用这个模型。序列化与反序列的目的是将数据、模型长久的保存。

2.PyTorch中的模型保存与加载

  PyTorch存储模型主要采用pkl，pt，pth三种格式。就使用层面来说没有区别。一个PyTorch模型主要包含两个部分：模型结构和权重。其中模型是继承nn.Module的类，权重的数据结构是一个字典（key是层名，value是权重向量）。因此，模型的保存与加载分别有两种模式：保存整个模型与保存模型参数。

2.1 模型保存

• torch.save
  模型保存（序列化）
  主要参数：
  • obj：对象（模型、张量、parameters、dict等）
  • f：输出路径（硬盘当中的路径，用于保存）

模式一：保存整个Module

torch.save(net,path)

模式二：保存模型参数（官方推荐）

state_dict=net.state_dict()
torch.save(state_dict,path)

  下面我们通过代码来简单观察模型保存两种模式的不同

import torch
import numpy as np
import torch.nn as nn


class LeNet2(nn.Module):
    def __init__(self, classes):
        super(LeNet2, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)
        return x

    def initialize(self):
        for p in self.parameters():
            p.data.fill_(20191104)


net = LeNet2(classes=10)

# "模拟训练"
# 查看第一个卷积层的第一个卷积核的参数
print("训练前: ", net.features[0].weight[0, ...])
net.initialize()
print("训练后: ", net.features[0].weight[0, ...])

path_model = "./model.pkl"
path_state_dict = "./model_state_dict.pkl"

# 保存整个模型
torch.save(net, path_model)

# 保存模型参数
net_state_dict = net.state_dict()
torch.save(net_state_dict, path_state_dict)

在当前文件目录下，就会生成model.pkl与model_state_dict.pkl。

2.2 模型加载

• torch.load
  模型加载（反序列化）
  主要参数：
  • f：文件路径
  • map_location：指定存放位置，cpu or gpu

模式一：加载整个Module

torch.load(net,path)

模式二：加载模型参数（官方推荐）

state_dict_load = torch.load(path_state_dict)
# 初始化模型
net_new = LeNet2(classes=10)
# 模型加载参数
net_new.load_state_dict(state_dict_load)

  下面我们通过代码来简单观察模型加载两种模式的不同

import torch
import numpy as np
import torch.nn as nn


class LeNet2(nn.Module):
    def __init__(self, classes):
        super(LeNet2, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)
        return x

    def initialize(self):
        for p in self.parameters():
            p.data.fill_(20191104)


# ================================== load net ===========================

path_model = "./model.pkl"
net_load = torch.load(path_model)
# 打印模型结构
print(net_load)

# ================================== load state_dict ===========================

path_state_dict = "./model_state_dict.pkl"
state_dict_load = torch.load(path_state_dict)

print(state_dict_load.keys())

# ================================== update state_dict ===========================

net_new = LeNet2(classes=10)

print("加载前: ", net_new.features[0].weight[0, ...])
net_new.load_state_dict(state_dict_load)
print("加载后: ", net_new.features[0].weight[0, ...])

LeNet2(
  (features): Sequential(
    (0): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
    (4): ReLU()
    (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (classifier): Sequential(
    (0): Linear(in_features=400, out_features=120, bias=True)
    (1): ReLU()
    (2): Linear(in_features=120, out_features=84, bias=True)
    (3): ReLU()
    (4): Linear(in_features=84, out_features=2019, bias=True)
  )
)

# 这些keys中的值全为20191104
odict_keys(['features.0.weight', 'features.0.bias', 'features.3.weight', 'features.3.bias', 'classifier.0.weight', 'classifier.0.bias', 'classifier.2.weight', 'classifier.2.bias', 'classifier.4.weight', 'classifier.4.bias'])

加载前:  tensor([[[ 0.1091,  0.0194, -0.0126, -0.0682,  0.1105],
         [-0.0846, -0.0898, -0.1119,  0.0616, -0.1045],
         [-0.0696,  0.0017, -0.0754,  0.0078, -0.0538],
         [-0.0208, -0.0018,  0.0816, -0.0684,  0.0662],
         [ 0.0453,  0.0400, -0.0184, -0.0038, -0.0663]],
        [[ 0.0226,  0.0914,  0.0725, -0.0074,  0.0676],
         [ 0.0181, -0.0345,  0.0596, -0.0887, -0.0850],
         [ 0.0062,  0.0234,  0.0859,  0.0583,  0.1044],
         [-0.1049,  0.0995, -0.0015,  0.0407,  0.0565],
         [ 0.0461, -0.0932,  0.0327, -0.0590, -0.0506]],
        [[-0.0758,  0.1004,  0.0691, -0.1083,  0.0960],
         [-0.1005, -0.0571,  0.0786,  0.0127,  0.0799],
         [-0.0487, -0.1015, -0.0787,  0.1073, -0.1136],
         [-0.0696,  0.0429, -0.0608,  0.0483, -0.1019],
         [-0.0581,  0.0846, -0.0643,  0.1144, -0.1064]]],
       grad_fn=<SelectBackward>)
加载后:  tensor([[[20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.]],
        [[20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.]],
        [[20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.],
         [20191104., 20191104., 20191104., 20191104., 20191104.]]],
       grad_fn=<SelectBackward>)

3.断点续训练

  模型的保存与加载不仅可以使得模型永久使用，而且在模型训练过程中也是十分有用的。常用技巧是断点续训练。断点续训练可以解决因某种意外的原因导致模型训练的终止，而需要重新训练，重复训练的问题。所以，在模型训练过程就要有一个断点续训练的机制，保存模型参数，以备在中断之后接着这个checkpoint继续训练。那么，断点续训练需要保存那些数据呢？在模型迭代训练过程中，模型与优化器不断变化着。当然，也可以保存loss，用来指示当前模型的状态。

# 基本参数
checkpoint = {
	"model_state_dict": net.state_dict(),
	"optimizer_state_dict": optimizer.state_dict(),
	"epoch": epoch
	}

  下面我们模拟断点续训练。

rmb_label = {"1": 0, "100": 1}

# 参数设置
checkpoint_interval = 5
MAX_EPOCH = 10
BATCH_SIZE = 16
LR = 0.01
log_interval = 10
val_interval = 1


# ============================ step 1/5 数据 ============================

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
split_dir = os.path.abspath(os.path.join(BASE_DIR, "..", "..", "materials", "rmb_split"))
train_dir = os.path.join(split_dir, "train")
valid_dir = os.path.join(split_dir, "valid")

if not os.path.exists(split_dir):
    raise Exception(r"数据 {} 不存在, 回到lesson-06\1_split_dataset.py生成数据".format(split_dir))

norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]

train_transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.RandomCrop(32, padding=4),
    transforms.RandomGrayscale(p=0.8),
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std),
])

valid_transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std),
])

# 构建MyDataset实例
train_data = RMBDataset(data_dir=train_dir, transform=train_transform)
valid_data = RMBDataset(data_dir=valid_dir, transform=valid_transform)

# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE)

# ============================ step 2/5 模型 ============================

net = LeNet(classes=2)
net.initialize_weights()

# ============================ step 3/5 损失函数 ============================
criterion = nn.CrossEntropyLoss()                                                   # 选择损失函数

# ============================ step 4/5 优化器 ============================
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9)                        # 选择优化器
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=6, gamma=0.1)     # 设置学习率下降策略

# ============================ step 5/5 训练 ============================
train_curve = list()
valid_curve = list()

start_epoch = -1
for epoch in range(start_epoch+1, MAX_EPOCH):

    loss_mean = 0.
    correct = 0.
    total = 0.

    net.train()
    for i, data in enumerate(train_loader):

        # forward
        inputs, labels = data
        outputs = net(inputs)

        # backward
        optimizer.zero_grad()
        loss = criterion(outputs, labels)
        loss.backward()

        # update weights
        optimizer.step()

        # 统计分类情况
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).squeeze().sum().numpy()

        # 打印训练信息
        loss_mean += loss.item()
        train_curve.append(loss.item())
        if (i+1) % log_interval == 0:
            loss_mean = loss_mean / log_interval
            print("Training:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
                epoch, MAX_EPOCH, i+1, len(train_loader), loss_mean, correct / total))
            loss_mean = 0.

    scheduler.step()  # 更新学习率

    if (epoch+1) % checkpoint_interval == 0:

        checkpoint = {"model_state_dict": net.state_dict(),
                      "optimizer_state_dict": optimizer.state_dict(),
                      "epoch": epoch}
        path_checkpoint = "./checkpoint_{}_epoch.pkl".format(epoch)
        torch.save(checkpoint, path_checkpoint)

    if epoch > 5:
        print("训练意外中断...")
        break

    # validate the model
    if (epoch+1) % val_interval == 0:

        correct_val = 0.
        total_val = 0.
        loss_val = 0.
        net.eval()
        with torch.no_grad():
            for j, data in enumerate(valid_loader):
                inputs, labels = data
                outputs = net(inputs)
                loss = criterion(outputs, labels)

                _, predicted = torch.max(outputs.data, 1)
                total_val += labels.size(0)
                correct_val += (predicted == labels).squeeze().sum().numpy()

                loss_val += loss.item()

            valid_curve.append(loss.item())
            print("Valid:\t Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
                epoch, MAX_EPOCH, j+1, len(valid_loader), loss_val/len(valid_loader), correct / total))


train_x = range(len(train_curve))
train_y = train_curve

train_iters = len(train_loader)
valid_x = np.arange(1, len(valid_curve)+1) * train_iters*val_interval # 由于valid中记录的是epochloss，需要对记录点进行转换到iterations
valid_y = valid_curve

plt.plot(train_x, train_y, label='Train')
plt.plot(valid_x, valid_y, label='Valid')

plt.legend(loc='upper right')
plt.ylabel('loss value')
plt.xlabel('Iteration')
plt.show()

Training:Epoch[000/010] Iteration[010/013] Loss: 0.6543 Acc:58.13%
Valid:	 Epoch[000/010] Iteration[003/003] Loss: 0.3089 Acc:64.62%
Training:Epoch[001/010] Iteration[010/013] Loss: 0.2285 Acc:94.38%
Valid:	 Epoch[001/010] Iteration[003/003] Loss: 0.0121 Acc:95.38%
Training:Epoch[002/010] Iteration[010/013] Loss: 0.1686 Acc:96.25%
Valid:	 Epoch[002/010] Iteration[003/003] Loss: 0.0582 Acc:95.38%
Training:Epoch[003/010] Iteration[010/013] Loss: 0.2316 Acc:90.00%
Valid:	 Epoch[003/010] Iteration[003/003] Loss: 0.0753 Acc:89.74%
Training:Epoch[004/010] Iteration[010/013] Loss: 0.1923 Acc:91.88%
Valid:	 Epoch[004/010] Iteration[003/003] Loss: 0.0067 Acc:93.33%
Training:Epoch[005/010] Iteration[010/013] Loss: 0.0573 Acc:98.75%
Valid:	 Epoch[005/010] Iteration[003/003] Loss: 0.0062 Acc:98.97%
Training:Epoch[006/010] Iteration[010/013] Loss: 0.0070 Acc:99.38%
训练意外中断...

在第5个epoch设置断点，接下来进行断点续训练

rmb_label = {"1": 0, "100": 1}

# 参数设置
checkpoint_interval = 5
MAX_EPOCH = 10
BATCH_SIZE = 16
LR = 0.01
log_interval = 10
val_interval = 1


# ============================ step 1/5 数据 ============================

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
split_dir = os.path.abspath(os.path.join(BASE_DIR, "..", "..", "materials", "rmb_split"))
train_dir = os.path.join(split_dir, "train")
valid_dir = os.path.join(split_dir, "valid")

if not os.path.exists(split_dir):
    raise Exception(r"数据 {} 不存在, 回到lesson-06\1_split_dataset.py生成数据".format(split_dir))

norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]

train_transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.RandomCrop(32, padding=4),
    transforms.RandomGrayscale(p=0.8),
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std),
])

valid_transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std),
])

# 构建MyDataset实例
train_data = RMBDataset(data_dir=train_dir, transform=train_transform)
valid_data = RMBDataset(data_dir=valid_dir, transform=valid_transform)

# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE)

# ============================ step 2/5 模型 ============================

net = LeNet(classes=2)
net.initialize_weights()

# ============================ step 3/5 损失函数 ============================
criterion = nn.CrossEntropyLoss()                                                   # 选择损失函数

# ============================ step 4/5 优化器 ============================
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9)                        # 选择优化器
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=6, gamma=0.1)     # 设置学习率下降策略


# ============================ step 5+/5 断点恢复 ============================
path_checkpoint = "./checkpoint_4_epoch.pkl"
checkpoint = torch.load(path_checkpoint)

net.load_state_dict(checkpoint['model_state_dict'])

optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

start_epoch = checkpoint['epoch']

scheduler.last_epoch = start_epoch

# ============================ step 5/5 训练 ============================
train_curve = list()
valid_curve = list()

for epoch in range(start_epoch + 1, MAX_EPOCH):

    loss_mean = 0.
    correct = 0.
    total = 0.

    net.train()
    for i, data in enumerate(train_loader):

        # forward
        inputs, labels = data
        outputs = net(inputs)

        # backward
        optimizer.zero_grad()
        loss = criterion(outputs, labels)
        loss.backward()

        # update weights
        optimizer.step()

        # 统计分类情况
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).squeeze().sum().numpy()

        # 打印训练信息
        loss_mean += loss.item()
        train_curve.append(loss.item())
        if (i+1) % log_interval == 0:
            loss_mean = loss_mean / log_interval
            print("Training:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
                epoch, MAX_EPOCH, i+1, len(train_loader), loss_mean, correct / total))
            loss_mean = 0.

    scheduler.step()  # 更新学习率

    if (epoch+1) % checkpoint_interval == 0:

        checkpoint = {"model_state_dict": net.state_dict(),
                      "optimizer_state_dic": optimizer.state_dict(),
                      "loss": loss,
                      "epoch": epoch}
        path_checkpoint = "./checkpint_{}_epoch.pkl".format(epoch)
        torch.save(checkpoint, path_checkpoint)

    # if epoch > 5:
    #     print("训练意外中断...")
    #     break

    # validate the model
    if (epoch+1) % val_interval == 0:

        correct_val = 0.
        total_val = 0.
        loss_val = 0.
        net.eval()
        with torch.no_grad():
            for j, data in enumerate(valid_loader):
                inputs, labels = data
                outputs = net(inputs)
                loss = criterion(outputs, labels)

                _, predicted = torch.max(outputs.data, 1)
                total_val += labels.size(0)
                correct_val += (predicted == labels).squeeze().sum().numpy()

                loss_val += loss.item()

            valid_curve.append(loss.item())
            print("Valid:\t Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
                epoch, MAX_EPOCH, j+1, len(valid_loader), loss_val/len(valid_loader), correct / total))


train_x = range(len(train_curve))
train_y = train_curve

train_iters = len(train_loader)
valid_x = np.arange(1, len(valid_curve)+1) * train_iters*val_interval # 由于valid中记录的是epochloss，需要对记录点进行转换到iterations
valid_y = valid_curve

plt.plot(train_x, train_y, label='Train')
plt.plot(valid_x, valid_y, label='Valid')

plt.legend(loc='upper right')
plt.ylabel('loss value')
plt.xlabel('Iteration')
plt.show()

发现：模型从第5轮接着开始训练

Training:Epoch[005/010] Iteration[010/013] Loss: 0.0603 Acc:97.50%
Valid:	 Epoch[005/010] Iteration[003/003] Loss: 0.0031 Acc:97.95%
Training:Epoch[006/010] Iteration[010/013] Loss: 0.0151 Acc:99.38%
Valid:	 Epoch[006/010] Iteration[003/003] Loss: 0.0021 Acc:99.49%
Training:Epoch[007/010] Iteration[010/013] Loss: 0.1177 Acc:97.50%
Valid:	 Epoch[007/010] Iteration[003/003] Loss: 0.0020 Acc:97.95%
Training:Epoch[008/010] Iteration[010/013] Loss: 0.0032 Acc:100.00%
Valid:	 Epoch[008/010] Iteration[003/003] Loss: 0.0019 Acc:98.97%
Training:Epoch[009/010] Iteration[010/013] Loss: 0.0160 Acc:99.38%
Valid:	 Epoch[009/010] Iteration[003/003] Loss: 0.0017 Acc:98.97%

4.单卡和多卡模型存储的区别

  PyTorch中将模型和数据放到GPU上有两种方式——.cuda()和.to(device)，如果要使用多卡训练的话，需要对模型使用torch.nn.DataParallel。示例如下：

os.environ['CUDA_VISIBLE_DEVICES'] = '0' # 如果是多卡改成类似0,1,2
model = model.cuda()  # 单卡
model = torch.nn.DataParallel(model).cuda()  # 多卡

之后我们把model对应的layer名称打印出来看一下，可以观察到差别在于多卡并行的模型每层的名称前多了一个“module”。

• 单卡模型的层名:
  
• 多卡模型的层名：
  

这种模型表示的不同可能会导致模型保存和加载过程中需要处理一些矛盾点，下面对各种可能的情况做分类讨论。

5.情况分类讨论

  由于训练和测试所使用的硬件条件不同，在模型的保存和加载过程中可能因为单GPU和多GPU环境的不同带来模型不匹配等问题。这里对PyTorch框架下单卡/多卡下模型的保存和加载问题进行排列组合。

• 单卡保存+单卡加载
  在使用os.envision命令指定使用的GPU后，即可进行模型保存和读取操作。注意这里即便保存和读取时使用的GPU不同也无妨。

  import os
  import torch
  from torchvision import models
  
  os.environ['CUDA_VISIBLE_DEVICES'] = '0'   #这里替换成希望使用的GPU编号
  model = models.resnet152(pretrained=True)
  model.cuda()
  
  # 保存+读取整个模型
  torch.save(model, save_dir)
  loaded_model = torch.load(save_dir)
  loaded_model.cuda()
  
  # 保存+读取模型权重
  torch.save(model.state_dict(), save_dir)
  loaded_dict = torch.load(save_dir)
  loaded_model = models.resnet152()   #注意这里需要对模型结构有定义
  loaded_model.state_dict = loaded_dict
  loaded_model.cuda()
  

• 单卡保存+多卡加载
  这种情况的处理比较简单，读取单卡保存的模型后，使用nn.DataParallel函数进行分布式训练设置即可。

  import os
  import torch
  from torchvision import models
  
  os.environ['CUDA_VISIBLE_DEVICES'] = '0'   #这里替换成希望使用的GPU编号
  model = models.resnet152(pretrained=True)
  model.cuda()
  
  # 保存+读取整个模型
  torch.save(model, save_dir)
  
  os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'   #这里替换成希望使用的GPU编号
  loaded_model = torch.load(save_dir)
  loaded_model = nn.DataParallel(loaded_model).cuda()
  
  # 保存+读取模型权重
  torch.save(model.state_dict(), save_dir)
  
  os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'   #这里替换成希望使用的GPU编号
  loaded_dict = torch.load(save_dir)
  loaded_model = models.resnet152()   #注意这里需要对模型结构有定义
  loaded_model.state_dict = loaded_dict
  loaded_model = nn.DataParallel(loaded_model).cuda()
  

• 多卡保存+单卡加载
  这种情况下的核心问题是：如何去掉权重字典键名中的"module"，以保证模型的统一性。

  • 对于加载整个模型，直接提取模型的module属性即可：

    import os
    import torch
    from torchvision import models
    
    os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'   #这里替换成希望使用的GPU编号
    
    model = models.resnet152(pretrained=True)
    model = nn.DataParallel(model).cuda()
    
    # 保存+读取整个模型
    torch.save(model, save_dir)
    
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'   #这里替换成希望使用的GPU编号
    loaded_model = torch.load(save_dir)
    loaded_model = loaded_model.module
    

  • 对于加载模型权重，往model里添加module（推荐）

    import os
    import torch
    from torchvision import models
    
    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'   #这里替换成希望使用的GPU编号
    
    model = models.resnet152(pretrained=True)
    model = nn.DataParallel(model).cuda()
    
    # 保存+读取模型权重
    torch.save(model.state_dict(), save_dir)
    
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'   #这里替换成希望使用的GPU编号
    loaded_dict = torch.load(save_dir)
    loaded_model = models.resnet152()   #注意这里需要对模型结构有定义
    loaded_model = nn.DataParallel(loaded_model).cuda()
    loaded_model.state_dict = loaded_dict
    

    这样即便是单卡，也可以开始训练了（相当于分布到单卡上）。也可以遍历字典去除module

    from collections import OrderedDict
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'   #这里替换成希望使用的GPU编号
    
    loaded_dict = torch.load(save_dir)
    
    new_state_dict = OrderedDict()
    for k, v in loaded_dict.items():
        name = k[7:] # module字段在最前面，从第7个字符开始就可以去掉module
        new_state_dict[name] = v #新字典的key值对应的value一一对应
    
    loaded_model = models.resnet152()   #注意这里需要对模型结构有定义
    loaded_model.state_dict = new_state_dict
    loaded_model = loaded_model.cuda()
    

    也可以使用replace操作去除module。

    loaded_model = models.resnet152()    
    loaded_dict = torch.load(save_dir)
    loaded_model.load_state_dict({k.replace('module.', ''): v for k, v in loaded_dict.items()})
    

• 多卡保存+多卡加载
  由于是模型保存和加载都使用的是多卡，因此不存在模型层名前缀不同的问题。但多卡状态下存在一个device（使用的GPU）匹配的问题，即保存整个模型时会同时保存所使用的GPU id等信息，读取时若这些信息和当前使用的GPU信息不符则可能会报错或者程序不按预定状态运行。具体表现为以下两点：

  • 读取整个模型再使用nn.DataParallel进行分布式训练设置
    这种情况很可能会造成保存的整个模型中GPU id和读取环境下设置的GPU id不符，训练时数据所在device和模型所在device不一致而报错。
  • 读取整个模型而不使用nn.DataParallel进行分布式训练设置
    这种情况可能不会报错，测试中发现程序会自动使用设备的前n个GPU进行训练（n是保存的模型使用的GPU个数）。此时如果指定的GPU个数少于n，则会报错。在这种情况下，只有保存模型时环境的device id和读取模型时环境的device id一致，程序才会按照预期在指定的GPU上进行分布式训练。

  相比之下，读取模型权重，之后再使用nn.DataParallel进行分布式训练设置则没有问题。因此多卡模式下建议使用权重的方式存储和读取模型：

  import os
  import torch
  from torchvision import models
  
  os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'   #这里替换成希望使用的GPU编号
  
  model = models.resnet152(pretrained=True)
  model = nn.DataParallel(model).cuda()
  
  # 保存+读取模型权重，强烈建议！！
  torch.save(model.state_dict(), save_dir)
  loaded_dict = torch.load(save_dir)
  loaded_model = models.resnet152()   #注意这里需要对模型结构有定义
  loaded_model = nn.DataParallel(loaded_model).cuda()
  loaded_model.state_dict = loaded_dict
  

  如果只有保存的整个模型，也可以采用提取权重的方式构建新的模型：

  # 读取整个模型
  loaded_whole_model = torch.load(save_dir)
  loaded_model = models.resnet152()   #注意这里需要对模型结构有定义
  loaded_model.state_dict = loaded_whole_model.state_dict
  loaded_model = nn.DataParallel(loaded_model).cuda()
  

  另外，上面所有对于loaded_model修改权重字典的形式都是通过赋值来实现的，在PyTorch中还可以通过"load_state_dict"函数来实现：

  loaded_model.load_state_dict(loaded_dict)
  

---

• 第二章、第三章、第四章部分内容参考于thorough-pytorch