PyTorch定义模型的三种方式

Pytorch定义模型的三种方式

Pytorch通常有三种方式构建模型：使用nn.Sequential按层顺序构建模型，继承nn.Module基类构建自定义模型，继承nn.Module基类构建模型并辅助应用模型容器(nn.Sequential,nn.ModuleList,nn.ModuleDict)进行封装。

第一种：用nn.Sequential按层顺序构建模型

def creatModel():
    model=nn.Sequential()
    model.add_module('conv1',nn.Conv2d(in_channels=3,out_channels=32,kernel_size=3))
    model.add_module('pool',nn.MaxPool2d(kernel_size=2,stride=2))
    model.add_module('conv2',nn.Conv2d(in_channels=32,out_channels=64,kernel_size=5))

    model.add_module('dropout',nn.Dropout2d(p=0.1))
    model.add_module('adaptive_pool',nn.AdaptiveMaxPool2d((1,1)))
    model.add_module('flatten',nn.Flatten())
    model.add_module('linear1',nn.Linear(64,32))
    model.add_module('relu',nn.ReLU())
    model.add_module('linear2',nn.Linear(32,1))
    model.add_module('sigmoid',nn.Sigmoid())
    return model

mymodel=creatModel()
print(mymodel)
torchkeras.summary(mymodel,input_shape=(3,512,512))

结果为：

Sequential(
  (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv2): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1))
  (dropout): Dropout2d(p=0.1, inplace=False)
  (adaptive_pool): AdaptiveMaxPool2d(output_size=(1, 1))
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear1): Linear(in_features=64, out_features=32, bias=True)
  (relu): ReLU()
  (linear2): Linear(in_features=32, out_features=1, bias=True)
  (sigmoid): Sigmoid()
)
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 32, 510, 510]             896
         MaxPool2d-2         [-1, 32, 255, 255]               0
            Conv2d-3         [-1, 64, 251, 251]          51,264
         Dropout2d-4         [-1, 64, 251, 251]               0
 AdaptiveMaxPool2d-5             [-1, 64, 1, 1]               0
           Flatten-6                   [-1, 64]               0
            Linear-7                   [-1, 32]           2,080
              ReLU-8                   [-1, 32]               0
            Linear-9                    [-1, 1]              33
          Sigmoid-10                    [-1, 1]               0
================================================================
Total params: 54,273
Trainable params: 54,273
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 3.000000
Forward/backward pass size (MB): 140.902115
Params size (MB): 0.207035
Estimated Total Size (MB): 144.109150
----------------------------------------------------------------

第二种：继承nn.Module基类构建自定义模型

import torch
from torch import nn
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel,self).__init__()
        self.conv1=nn.Conv2d(in_channels=3,out_channels=32,kernel_size=3)
        self.pool=nn.MaxPool2d(kernel_size=2,stride=2)
        self.conv2=nn.Conv2d(in_channels=32,out_channels=64,kernel_size=5)
        self.dropout=nn.Dropout2d(p=0.1)
        self.adaptive_pool=nn.AdaptiveMaxPool2d((1,1))
        self.flatten=nn.Flatten()
        self.linear1=nn.Linear(64,32)
        self.relu=nn.ReLU()
        self.linear2=nn.Linear(32,1)
        self.sigmoid=nn.Sigmoid()

    def forward(self,x):
        x=self.conv1(x)
        x=self.pool(x)
        x=self.conv2(x)
        x=self.dropout(x)
        x=self.adaptive_pool(x)
        x=self.flatten(x)
        x=self.linear1(x)
        x=self.relu(x)
        x=self.linear2(x)
        y=self.sigmoid(x)
        return y
model =MyModel()
import torchkeras
torchkeras.summary(model,input_shape=(3,512,512))

运行结果如下：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 32, 510, 510]             896
         MaxPool2d-2         [-1, 32, 255, 255]               0
            Conv2d-3         [-1, 64, 251, 251]          51,264
         Dropout2d-4         [-1, 64, 251, 251]               0
 AdaptiveMaxPool2d-5             [-1, 64, 1, 1]               0
           Flatten-6                   [-1, 64]               0
            Linear-7                   [-1, 32]           2,080
              ReLU-8                   [-1, 32]               0
            Linear-9                    [-1, 1]              33
          Sigmoid-10                    [-1, 1]               0
================================================================
Total params: 54,273
Trainable params: 54,273
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 3.000000
Forward/backward pass size (MB): 140.902115
Params size (MB): 0.207035
Estimated Total Size (MB): 144.109150
----------------------------------------------------------------

第三种：继承nn.Module基类构建模型并辅助应用模型容器(nn.Sequential,nn.ModuleList,nn.ModuleDict)进行封装

当模型比较复杂时，可使用