深入浅出pytorch笔记——第五章

第五章

5.1PyTorch模型定义的方式

Sequential

1. 可以接收一个子模块的有序字典(OrderedDict) 或者一系列子模块作为参数来逐一添加 Module 的实例，而模型的前向计算就是将这些实例按添加的顺序逐⼀计算
2. 使用Sequential定义模型的好处在于简单、易读，使用Sequential定义的模型不需要再写forward，因为顺序已经定义好了
3. 使用Sequential也会使得模型定义丧失灵活性，比如需要在模型中间加入一个外部输入时就不适合用Sequential的方式实现。使用时需根据实际需求加以选择。

直接排列

import torch.nn as nn
net = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10), 
        )
print(net)

Sequential(
  (0): Linear(in_features=784, out_features=256, bias=True)
  (1): ReLU()
  (2): Linear(in_features=256, out_features=10, bias=True)
)

使用OrderedDict

1. 可以实现层命名

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)

Sequential(
  (fc1): Linear(in_features=784, out_features=256, bias=True)
  (relu1): ReLU()
  (fc2): Linear(in_features=256, out_features=10, bias=True)
)

ModuleList

1. ModuleList 接收一个子模块（或层，需属于nn.Module类）的列表作为输入，然后也可以类似List那样进行append和extend操作
2. 子模块或层的权重也会自动添加到网络中来

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

Linear(in_features=256, out_features=10, bias=True)
ModuleList(
  (0): Linear(in_features=784, out_features=256, bias=True)
  (1): ReLU()
  (2): Linear(in_features=256, out_features=10, bias=True)
)

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

class model(nn.Module):
  def __init__(self, ...):
    super().__init__()
    self.modulelist = ...
    ...
    
  def forward(self, x):
    for layer in self.modulelist:
      x = layer(x)
    return x

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)

Linear(in_features=784, out_features=256, bias=True)
Linear(in_features=256, out_features=10, bias=True)
ModuleDict(
  (act): ReLU()
  (linear): Linear(in_features=784, out_features=256, bias=True)
  (output): Linear(in_features=256, out_features=10, bias=True)
)

ModuleDict并没有定义一个网络，它只是将不同的模块储存在一起。需要经过forward函数指定各个层的先后顺序后才算完成了模型的定义，具体实现时用for循环即可完成。

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

U-Net分析

组成U-Net的模型块主要有如下几个部分：
1）每个子块内部的两次卷积（Double Convolution）
2）左侧模型块之间的下采样连接，即最大池化（Max pooling）
3）右侧模型块之间的上采样连接（Up sampling）
4）输出层的处理
除模型块外，还有模型块之间的横向连接，输入和U-Net底部的连接等计算，这些单独的操作可以通过forward函数来实现

模块化

Double Convolution

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""
    #因为有两次卷积， mid_channels就是第一个卷积的输出和第二个卷积的输入
    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)

Down Sampling

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),
            #引用了Double Convolution
            DoubleConv(in_channels, out_channels)
        )
    def forward(self, x):
        return self.maxpool_conv(x)

Up Sampling(包含skip connection)

class Up(nn.Module):
    """Upscaling then double conv"""
    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()
        #Up Sampling是一个上采样的过程，需要插值，bilinear是插值方式
        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            #使用pytorch内置函数进行上采样操作
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            #//是地板除，向下取整
            #引用Double Convolution
            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)
    #包含skip connection
    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)

OutConv

class OutConv(nn.Module):
    #out_channels取决于你是几分类问题
    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)

组装

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        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

5.3 PyTorch修改模型

修改模型层

resnet50模型结构是为了适配ImageNet预训练的权重，因此最后全连接层（fc）的输出节点数是1000
假设我们要用这个resnet模型去做一个10分类的问题，就应该修改模型的fc层，将其输出节点数替换为10

import torchvision.models as models
net = models.resnet50()
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

添加外部输入

将原模型添加输入位置前的部分作为一个整体，同时在forward中定义好原模型不变的部分、添加的输入和后续层之间的连接关系，从而完成模型的修改。
在倒数第二层增加一个额外的输入变量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):
        #这里是上述的resnet50
        x = self.net(x)
        #倒数第二层增加一个额外的输入变量，通过torch.cat实现了tensor的拼接
        #对外部输入变量"add_variable"进行unsqueeze操作是为了和net输出的tensor保持维度一致，从而可以和tensor进行torch.cat操作
        x = torch.cat((self.dropout(self.relu(x)), add_variable.unsqueeze(1)),1)
        x = self.fc_add(x)
        x = self.output(x)
        return x

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

outputs = model(inputs, add_var)

添加额外输出

需要输出模型某一中间层的结果，以施加额外的监督，获得更好的中间层结果。修改模型定义中forward函数的return变量

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

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

out10, out1000 = model(inputs, add_var)

5.4 PyTorch模型保存与读取

PyTorch中将模型和数据放到GPU上有两种方式——.cuda()和.to(device)，我们针对.cuda()讨论

模型存储内容

一个PyTorch模型主要包含两个部分：模型结构和权重

1. 模型是继承nn.Module的类
2. 权重的数据结构是一个字典，key是层名，value是权重向量
   存储分为两种形式
3. 存储整个模型（包括结构和权重）
4. 只存储模型权重。
   对于PyTorch而言，pt, pth和pkl三种数据格式均支持模型权重和整个模型的存储，因此使用上没有差别

保存&读取

分单卡和多卡

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

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

CPU&单卡保存

一般建议保存整个模型，存在于cuda0

整个模型

#保存
torch.save(mm, "save_dir.model.pth")

#单卡读取
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
load_model=torch.load("save_dir.model.pth")
#多卡读取
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'  
load_model=torch.load("save_dir.model.pth")

#单卡实例化新模型
loaded_model.cuda()
#多卡实例化新模型
loaded_model = nn.DataParallel(loaded_model).cuda()

#通过实例化的模型读取权重
load_model.state_dict()

模型权重

#保存
torch.save(mm.state_dict(), "save_dir.weights.pth")

#单卡读取
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
load_weights=torch.load("save_dir.weights.pth")
#多卡读取
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'  
load_weights=torch.load("save_dir.weights.pth")

#1通过赋值来实现修改权重字典，把权重赋给新模型
mm_xin.state_dict = load_weights
#或2通过load_state_dict"函数来实现
mm_xin.load_state_dict(load_weights)

#单卡实例化新模型
mm_xin.cuda()
#多卡实例化新模型
mm_xin = nn.DataParallel(loaded_model).cuda()

#通过实例化的模型读取权重
mm_xin.state_dict()

多卡保存，多卡读取

多卡时模型分布在多个GPU上，即cuda1，cuda2，cuda3等等，模型加载时如果换机器（即GPU数量不匹配）的话，会很难加载成功（模型每层的名称前的module导致的）

保存整个模型，读取时采用提取权重的方式构建新模型

#保存
torch.save(mm_mul, "save_dir.model.pth")

#读取保存的整个模型的权重，给我们的新模型
mm_mul_xin.state_dict=mm_mul.state_dict

#利用提取的权重构建新模型
mm_mul_xin= torch.nn.DataParallel(mm_mul_xin).cuda()  # 多卡

#读取
mm_mul_xin.state_dict()

模型权重

#保存
torch.save(mm_mul.state_dict(), "save_dir.weights.pth")

#读取
load_weights_mul=torch.load("save_dir.weights.pth")

#通过实例化的模型，进行load_state_dict，加载刚刚读取的变量
mm_mul.load_state_dict(load_weights_mul)

#通过实例化的模型读取权重
mm_mul.state_dict()

多卡保存，单卡读取

这种情况下的核心问题是：如何去掉权重字典键名中的"module"，以保证模型的统一性

整个模型

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

#保存
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)
##取模型的module属性
loaded_model = loaded_model.module

模型权重

往model里添加module

去除字典里的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()})

STUffT&夏日回音