Linux深入浅出PyTorch（三）PyTorch模型定义

PyTorch模型定义的方式

模型是深度学习中重要的一部分。深度学习中有很多主流模型：

• CNN：解决了图像、视频处理中的诸多问题；
• RNN/LSTM：解决了序列数据处理的问题；
• GNN：在图模型上发挥着重要的作用；等等

1. nn.Module 类回顾

• torch.nn 模块里提供的一个模型构造类，是所有神经网络模块的基类，可以通过继承它来定义想要的模型；
• PyTorch模型定义应包括两个主要部分：构造函数__init__（用于创建模型参数）和构造函数forward（定义前向计算）
  注：如果仅仅定义了__init__ 和forward 函数，而没有继承nn.Module类，这样仅仅是把层定义了出来，也有了数据流向，也就是能正常进行前向计算，但是无法进行反向传播计算。[2]
• 基于nn.Module，可以通过：Sequential，Modulelist，ModuleDict 三种方式定义PyTorch模型。

2. Sequential

当一个模型较简单的时候，我们可以使用torch.nn.Sequential类来实现简单的顺序连接模型。它可以接收一个子模块的有序字典(OrderedDict) 或者一系列子模块作为参数来逐一添加 Module 的实例，⽽模型的前向计算就是将这些实例按添加的顺序逐⼀计算。结合Sequential 的定义方式理解如下：

class MySequential(nn.Module):
    from collections import OrderedDict
    def __init__(self, *args):
        super(MySequential, self).__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict): # 如果传入的是一个OrderedDict
            for key, module in args[0].items():
                self.add_module(key, module)  
                # add_module方法会将module添加进self._modules(一个OrderedDict)
        else:  # 传入的是一些Module
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)
    def forward(self, input):
        # self._modules返回一个 OrderedDict，保证会按照成员添加时的顺序遍历成
        for module in self._modules.values():
            input = module(input)
        return input

那么如何使用Sequential 来定义模型呢？有以下三种方式[1][3]：
（1）直接排列
如下：

import torch.nn as nn
net = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10), 
        )
print(net)
print(net[2]) # 通过索引获取某一层

注：这样做有一个问题，每一个层是没有名称，默认的是以0、1、2来命名。

（2）使用OrderedDict

import collections
import torch.nn as nn
net2 = nn.Sequential(collections.OrderedDict([
          ('fc1', nn.Linear(784, 256)),
          ('relu1', nn.ReLU()),
          ('fc2', nn.Linear(256, 10))
          ]))
print(net2)
print(net2[2]) # 通过索引获取某一层

注：此时每个层都有了自己的名字，但是需要注意的是，并不能通过名称直接获取层，依然只能通过索引index，即：model[2] 是正确的，model[“fc2”] 是错误的。

（3）通过 .add_module() 逐个添加

import torch.nn as nn
from collections import OrderedDict
model = nn.Sequential()
model.add_module("conv1",nn.Conv2d(1,20,5))
model.add_module('relu1', nn.ReLU())
model.add_module('conv2', nn.Conv2d(20,64,5))
model.add_module('relu2', nn.ReLU())
 
print(model)
print(model[2]) # 通过索引获取某一层

如果
注：add_module() 方法定义在父类Module 里面，不在Sequential 中。

总结：

• 使用Sequential 定义的模型简单、易读；
• 使用Sequential定义的模型不需要再写forward函数，因为顺序已经定义好了；
• 使用Sequential也会使得模型定义丧失灵活性，如果需要在模型中间加入一个外部输入时就不适合用Sequential的方式实现。

3. ModuleList

• 接收一个子模块（或层，需属于nn.Module类）的列表作为输入；
• 可以类似List那样进行append和extend操作(但是如果在定义网络net的代码前面不加nn.ModuleList ，只是将他们写成list的形式，则在forward计算的时候是正常的，不能进行反向计算，网络参数也不会被更新，因为使用 Python 的 list 添加的卷积层和它们的 parameters 并没有自动注册到网络中[4])；
• 加入到 nn.ModuleList 里面的 module 是会自动注册到整个网络上的，同时 module 的 parameters 也会自动添加到整个网络中。[4]
  如：

net = nn.ModuleList([nn.Linear(784, 256), nn.ReLU()])
net.append(nn.Linear(256, 10)) # # 类似List的append操作
print(net[-1])  # 类似List的索引访问
print(net)

注：

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

class net_modlist(nn.Module):
    def __init__(self):
        super(net_modlist, self).__init__()
        self.modlist = nn.ModuleList([
                       nn.Conv2d(1, 20, 5),
                       nn.ReLU(),
                        nn.Conv2d(20, 64, 5),
                        nn.ReLU()
                        ])

    def forward(self, x):
        for m in self.modlist:
            x = m(x)
        return x

net_modlist = net_modlist()
print(net_modlist)

总结：nn.Sequential与nn.ModuleList的区别[4]：

1. nn.Sequential内部实现了forward函数，因此可以不用写forward函数。而nn.ModuleList则没有实现内部forward函数；
2. nn.Sequential可以使用OrderedDict对每层进行命名；
3. nn.Sequential里面的模块按照顺序进行排列的，所以必须确保前一个模块的输出大小和下一个模块的输入大小是一致的；而nn.ModuleList 并没有定义一个网络，它只是将不同的模块储存在一起，这些模块之间并没有什么先后顺序可言；
4. 有的时候网络中有很多相似或者重复的层，我们一般会考虑用 for 循环来创建它们，而不是一行一行地写，这里用nn.ModuleList 实现。

4. ModuleDict

ModuleDict和ModuleList的作用类似，只是ModuleDict能够更方便地为神经网络的层添加名称：

net = nn.ModuleDict({
    
    
    'linear': nn.Linear(784, 256),
    'act': nn.ReLU(),
})
net['output'] = nn.Linear(256, 10) # 添加
print(net['linear']) # 访问
print(net.output)
print(net)

注：nn.ModuleDict 可以通过名称索引对应网络层，nn.Sequential 不可以。

5. 三种方法的比较

三种方法的不同点及适用场景如下[5]：

• nn.Sequential ：按顺序包装多个网络层，注重顺序性，各网络层之间严格按照顺序执行，不需要同时写__init__和 forward；
• nn.moduleList ：像Python的 list 一样包装多个网络层，注重迭代性，常用于大量重复网络构建，通过for循环实现重复构建；
• nn.ModuleDict ：像Python的 dict 一样包装多个网络层，注重索引性，常用于可选择的网络层；
• 当我们需要之前层的信息的时候，比如 ResNets 中的残差计算，当前层的结果需要和之前层中的结果进行融合，一般使用 ModuleList/ModuleDict 比较方便。

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

• 当模型的深度非常大时候，使用Sequential定义模型结构需要向其中添加几百行代码，使用起来不甚方便；
• 将模型中重复出现的层结构定义成一个”模块“，每次只需要向网络中添加对应的模块来构建模型，这样将会极大便利模型的构建过程；
• 以下以U-Net为例，介绍如何构建模型块，以及如何利用模型块快速搭建复杂模型。

1. U-Net简介

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

2. U-Net模型分析

• 模型具有非常好的对称性(U-Net网络的名字也是取自其”U“形的形状)
• 模型从上到下分为若干层，每层由左侧和右侧两个模型块组成，每侧的模型块与其上下模型块之间有连接；
• 同时位于同一层左右两侧的模型块之间也有连接，称为“Skip-connection”；
• 还包括输入和输出处理等其他组成部分；
• 组成U-Net的模型块主要部分如下：
  • 每个子块内部的两次卷积（Double Convolution）
  • 左侧模型块之间的下采样连接，即最大池化（Max pooling）
  • 右侧模型块之间的上采样连接（Up sampling）
  • 输出层的处理
• 除模型块外，还有模型块之间的横向连接，输入和U-Net底部的连接等计算，这些单独的操作可以通过forward函数来实现。

3. U-Net模型块实现

组建大型网络时，不必把每一层按序排列显式写出，这样太麻烦且不宜读，一种比较好的方法是先定义好模型块，再定义模型块之间的连接顺序和计算方式。

这里对U-Net的实现先组建4个模型块：DoubleConv, Down, Up, OutConv，每个模型块的PyTorch实现如下：

# 导入必要的包
import torch
import torch.nn as nn
import torch.nn.functional as F

# 模型块1：DoubleConv
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)

# 模型块2：Down
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)

# 模型块3：Up
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])

        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

# 模型块4：OutConv
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修改模型

有时需要修改模型的结构和添加额外输入输出。

1. 修改模型层

以PyTorch官方视觉库torchvision预定义好的模型 ResNet50 为例，探索如何修改模型的某一层或者某几层。模型定义如下：

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

这里模型结构是为了适配 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)), # 第二层fc层，且设置输出为10
                          ('output', nn.Softmax(dim=1))
                          ]))
    
net.fc = classifier

注：这里的操作相当于将模型（net）最后名称为“fc”的层替换成了名称为“classifier”的结构，classifier结构是自己用Sequential+OrderedDict方式定义的，这样就完成了模型的修改，可以用于10分类任务了。

2.添加外部输入

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

class Model(nn.Module):
    def __init__(self, net):
        super(Model, self).__init__()
        self.net = net # ResNet原模型
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc_add = nn.Linear(1001, 10, bias=True) # 加一层fc层，并且这里输出变成了1000+1=1001（因为新加了1维输入）
        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) # 将原来的tensor与新加的1维输入通过torch.cat进行拼接
        x = self.fc_add(x)
        x = self.output(x)
        return x

注：

• 为了让原来的1000维tensor与新加的1维输入成功拼接，需要将新加的输入通过 .unsqueeze(1) 操作变成与原来的tensor一样的维度
• torch.cat： 在给定维度上对输入的张量序列seq 进行连接操作，格式如下[6]：

torch.cat(inputs, dim=) # 输出一个Tensor

参数：
inputs：待连接的张量序列，可以是任意相同Tensor类型的Python序列
dim：选择的扩维，必须在[0, len(inputs[0])]之间，沿着此维连接张量序列

• torch.unsqueeze()：对数据维度进行扩充。给指定位置加上维数为1的维度。且其返回张量与输入张量共享内存，改变其中一个的内容另一个也会改变[7]。

torch.unsqueeze(input, dim, out=None)

参数：
input：输入张量
dim：int型，插入维度的索引
out：Tensor型，optional，结果张量

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

import torchvision.models as models
net = models.resnet50() # 原ResNet50的网络
model = Model(net).cuda() # 更改后的网络

另外，训练中在输入数据的时候要给两个input：

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) # 加1层fc，且输出变成10
        self.output = nn.Softmax(dim=1)
        
    def forward(self, x):
        x1000 = self.net(x) # 原网络
        x10 = self.dropout(self.relu(x1000))
        x10 = self.fc1(x10) # 加1层的fc计算
        x10 = self.output(x10)
        return x10, x1000 # 打印模型最后的输出和原网络的输出值

之后对修改好的模型结构进行实例化：

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

训练结果是两个输出：

out10, out1000 = model(inputs)

PyTorch模型保存与读取

1. 模型存储格式

• PyTorch存储模型主要采用pkl，pt，pth三种格式。在格式上并没有区别，只是后缀不同而已（仅此而已）。

2. 模型存储内容

• 一个PyTorch模型主要包含两个部分：模型结构和权重；
• 模型是继承nn.Module的类，权重的数据结构是一个字典（key是层名，value是权重向量）；
• 存储也由此分为两种形式：存储整个模型（包括结构和权重），和只存储模型权重；
• 对于PyTorch而言，pt, pth和pkl三种数据格式均支持模型权重和整个模型的存储。

from torchvision import models
model = models.resnet152(pretrained=True)

# 保存整个模型
torch.save(model, save_dir)
# 保存模型权重
torch.save(model.state_dict, save_dir)

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

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()  # 多卡

注：多卡并行的模型每层的名称前比单卡模型名称多了一个“module”

4. 模型保存和读取情况分类讨论

由于训练和测试所使用的硬件条件不同，在模型的保存和加载过程中可能因为单GPU和多GPU环境的不同带来模型不匹配等问题。
讨论各种可能的情况如下，样例模型是torchvision中预训练模型ResNet152：

（1）单卡保存 + 单卡加载

在使用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()

（2）单卡保存 + 多卡加载

读取单卡保存的模型后，使用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()

（3）多卡保存 + 单卡加载

这种情况下的核心问题是：如何去掉权重字典键名中的"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 # 比单卡保存 + 单卡读取多加一步

对于加载模型权重，有以下几种思路：
1）去除字典里的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

这样即便是单卡，也可以开始训练了（相当于分布到单卡上）

2）遍历字典去除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()

3）使用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()})

（4）多卡保存 + 多卡加载

• 由于是模型保存和加载都使用的是多卡，因此不存在模型层名前缀不同的问题（即每个保存的模型层名字前面都有”module“）；
• 多卡状态下存在一个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)

引用：
[1] datawhalechina/thorough-pytorch：https://github.com/datawhalechina/thorough-pytorch（强推）
[2] pytorch中几种构建模型的方法的区别：https://zhuanlan.zhihu.com/p/144400422
[3] pytorch教程之nn.Sequential类详解——使用Sequential类来自定义顺序连接模型：https://www.cnblogs.com/jiangkejie/p/13050566.html
[4] 详解PyTorch中的ModuleList和Sequential：https://zhuanlan.zhihu.com/p/75206669
[5] 模型容器Containers之nn.ModuleDict：https://www.cnblogs.com/gczr/p/14674688.html
[6] torch.cat()函数的官方解释，详解以及例子：https://blog.csdn.net/xinjieyuan/article/details/105208352
[7] torch.unsqueeze() 和 torch.squeeze()：https://zhuanlan.zhihu.com/p/86763381