PyTorch的模型通过继承nn.Moudle类,在类的内部定义子模块的实例化,通过前向计算调用子模块,最后实现
深度学习模型的搭建。
基础结构:
import torch.nn as nn class MyNet(nn.Module): def __init__(self, ...): # 定义类的初始化函数,...是用户的传入参数 super(MyNet, self).__init__() # 调用父类的初始化方法 ... # 根据传入的参数来定义子模块 def forward(self, ...): # 前向计算 ret = ... # 根据传入的张量和子模块计算返回张量 return ret
整个模块的函数主要由两部分组成:通过_init_方法初始化整个模型,forward方法对该模型进行前向计算。
示例:
import torch import torch.nn as nn from torchvision.models import alexnet class AlexNet(nn.Module): def __init__(self, num_classes=1000): super(AlexNet, self).__init__() 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)) 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) ) def forward(self, x): x = self.features(x) x = self.avgpool(x) x = torch.flatten(x, 1) x = self.classifier(x) return x