PyTorch的模型定义

（一）模型定义方式

1.当需要快速验证结果时，可以利用sequential搭建网络模型，因为已经明确了要用哪些层，直接写书写下来即可，不需要写初始化__init__和前向传播forward；

import torch
import torch.nn as nn
net1=nn.Sequential(
    nn.Linear(512,256),
    nn.ReLU(),
    nn.Linear(256,10)
)
print(net1)

#利用ordered可以实现模型层前面的名称命名
import torch
import torch.nn as nn
import collections
net=nn.Sequential(collections.OrderedDict([
    ('fc1',nn.Linear(512,256)),
    ('relu1',nn.ReLU()),
    ('fc2',nn.Linear(256,10))
]))
print(net2)

#输出的模型结果
Sequential(
  (0): Linear(in_features=512, out_features=256, bias=True)
  (1): ReLU()
  (2): Linear(in_features=256, out_features=10, bias=True)
)

2.当模型中的某一层需要重复出现多次时，ModuleList 可以很方便的实现一项多行。ModuleList 接收一个子模块（或层，需属于nn.Module类）的列表作为输入，然后也可以类似list那样进append和extend操作。同时，子模块或层的权重也会自动添加到网络中来。可以理解为它是一个储存不同 module，并自动将每个 module 的 parameters 添加到网络之中的容器。

import torch
import torch.nn as nn
net3=nn.ModuleList([nn.Linear(512,256),nn.ReLU()])
net3.append(nn.Linear(256,10))
print(net3[-1]) #类似list的索引访问
print(net3)

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

class net1(nn.Module):
    def __init__(self):
        super(net1, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(256,10) for i in range(3)])#3个全连接层
    def forward(self, x):
        for m in self.linears:
            x = m(x)
        return x

3.ModuleDict其实跟sequential里的ordered很像，它能够很方便地为神经网络的层添加名称。

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

注：同样ModuleDict也并没有定义一个网络，它只是将不同的模块存储在一起。所以不能直接将数据输入进去，会报错。

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

       在我们遇到的复杂模型中，往往每个模型多多少少都会有很多重复出现的结构。考虑到每一层有其输入和输出，若干层串联成的”模块“也有其输入和输出，如果我们能将这些重复出现的层定义为一个”模块“，每次只需要向网络中添加对应的模块来构建最终的模型，这样将会极大便利和简洁模型构建的过程。

以U-Net模型为例：

#下采样
class Downsample(nn.Module):
    def __init__(self,in_channels,out_channels):
        super(Downsample, self).__init__()
        self.conv_relu=nn.Sequential(
                        nn.Conv2d(in_channels,out_channels,kernel_size=3,padding=1),
                        nn.ReLU(),
                        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
                        nn.ReLU(),
        )
        self.pool=nn.MaxPool2d(kernel_size=2)
    def forward(self,x,is_pool=True):
        if is_pool:
            x=self.pool(x)
        x=self.conv_relu(x)
        return x

#上采样
class upsample(nn.Module):
    def __init__(self,channels):
        super(upsample,self).__init__()
        self.conv_relu=nn.Sequential(
            nn.Conv2d(2*channels,channels,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
            nn.ReLU(),
        )
        self.upconv=nn.Sequential(
            nn.ConvTranspose2d(channels,channels//2,kernel_size=3,stride=2,padding=1,output_padding=1),
            nn.ReLU()
        )
    def forward(self, x):
        x = self.conv_relu(x)
        x=self.upconv(x)
        return x

class Unet_model(nn.Module):
    def __init__(self):
        super(Unet_model,self).__init__()
        self.down1=Downsample(3,64)
        self.down2=Downsample(64,128)
        self.down3=Downsample(128,256)
        self.down4 = Downsample(256, 512)
        self.down5 = Downsample(512, 1024)
        self.up=nn.Sequential(
            nn.ConvTranspose2d(1024,512,kernel_size=3,stride=2,padding=1,output_padding=1),
            nn.ReLU(),
        )
        self.up1=upsample(512)
        self.up2 = upsample(256)
        self.up3 = upsample(128)
        self.conv_2=Downsample(128,64)
        self.last=nn.Conv2d(64,2,kernel_size=1)
    def forward(self,input):
        x1=self.down1(input,is_pool=False)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        x5 = self.down5(x4)

        x5=self.up(x5)
        x5=torch.cat([x4,x5],dim=1)
        x5=self.up1(x5)
        x5 = torch.cat([x3, x5], dim=1)
        x5 = self.up2(x5)
        x5 = torch.cat([x2, x5], dim=1)
        x5 = self.up3(x5)
        x5 = torch.cat([x1, x5], dim=1)
        x5=self.conv_2(x5,is_pool=False)
        output=self.last(x5)
        return output

（三）模型修改

1.修改模型特定层

import torch.nn as nn
import torchvision
from collections import OrderedDict
model=torchvision.models.vgg16()
fc = nn.Sequential(OrderedDict([('fc1', nn.Linear(1024, 512)),
                                        ('relu1', nn.ReLU()),
                                        ('dropout1', nn.Dropout(0.5)),
                                        ('fc2', nn.Linear(512, 256)),
                                        ('relu2', nn.ReLU()),
                                        ('dropout2', nn.Dropout(0.5)),
                                        ('fc3', nn.Linear(256, 10))
                                        ]))

model.classifier = fc #这里是将模型最后名称为“classifier”的层替换成了名称为“fc”的结构
print(model)

2.添加额外的输入

       基本的思路就是将原模型添加输入位置前的部分作为一个整体，同时在forward中定义好原模型不变的部分和添加的输入与后续层之间的连接关系，从而完成模型的修改。

class Model(nn.Module):
    def __init__(self, net):
        super(Model, self).__init__()
        self.net = net
        self.fc_add = nn.Linear(257, 10, bias=True)

    def forward(self, x, add_variable):
        x = self.net(x)
        x = x+ add_variable # add_variable可以是常数，也可以是数组，只需记住在添加完后，上面初始化的通道数也要做相应的改变
        x = self.fc_add(x)
        return x

3.添加额外的输出

基本思路就是在前向传播forward函数的renturn变量中增加输出的变量。

class Model(nn.Module):
    def __init__(self, net):
        super(Model, self).__init__()
        self.net = net
        self.fc_add = nn.Linear(257, 10, bias=True)

    def forward(self, x):
        x1= self.net(x)
        x2= self.fc_add(x)
        return x1,x2

（四）模型保存与读取

       模型的保存主要包含两个部分：模型结构和权重。其中模型是继承nn.Module的类，权重的数据结构是一个字典（key是层名，value是权重向量）。存储也由此分为两种形式：存储整个模型（包括结构和权重），和只存储模型权重。格式一般为以下三种：pkl，pt，pth。本次仅讨论在单卡和CPU情况下的模型保存与读取。

import torch
import torch.nn as nn
import torchvision
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'   #这里选择你想使用的GPU
model=torchvision.models.vgg16(pretrained=True)
model.to('cuda')

#第一种：CPU或单卡：保存、读取整个模型
PATH=('vgg_model.pth')
torch.save(model,PATH)#保存模型
load_model=torch.load(PATH)   #将保存的模型装载进去

#第二种：CPU或单卡：保存、读取模型权重
PATH=('vgg_parament.pth')
torch.save(model.state_dict(),PATH)#保存模型可训练参数的权重
vgg=torchvision.models.vgg16(pretrained=False)#在装载模型参数时必须先对模型的结构有定义
vgg.load_state_dict(torch.load(PATH))

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