import torch import torch.nn as nn

定义 Resnet 结构

class ResNet(nn.Module): def init(self, input_channel, output_channel, kernel_size): super(ResNet, self).init() self.conv1 = nn.Conv2d(input_channel, output_channel, kernel_size, padding=1) self.bn1 = nn.BatchNorm2d(output_channel) self.relu1 = nn.ReLU(inplace=True) self.conv2 = nn.Conv2d(output_channel, output_channel, kernel_size, padding=1) self.bn2 = nn.BatchNorm2d(output_channel) self.relu2 = nn.ReLU(inplace=True) self.conv3 = nn.Conv2d(output_channel, output_channel, kernel_size, padding=1) self.bn3 = nn.BatchNorm2d(output_channel) self.relu3 = nn.ReLU(inplace=True) self.add = nn.Add() self.max_pool = nn.MaxPool2d(2, 2)

def forward(self, x):
    # 第一组
    x = self.conv1(x)
    x = self.bn1(x)
    x = self.relu1(x)
    # 第二组
    x = self.conv2(x)
    x = self.bn2(x)
    x = self.relu2(x)
    # 第三组
    x = self.conv3(x)
    x = self.bn3(x)
    x = self.relu3(x)
    # 残差块
    shortcut = x
    x = self.add(x, shortcut)
    x = self.max_pool(x)
    return x

定义视频处理模块

class VideoProcess(nn.Module): def init(self, input_channel, output_channel): super(VideoProcess, self).init() self.resnet = ResNet(input_channel, output_channel, 3) self.fc1 = nn.Linear(output_channel, 128) self.fc2 = nn.Linear(128, 2)

def forward(self, x):
    x = self.resnet(x)
    x = x.view(x.size(0), -1)
    x = self.fc1(x)
    x = self.fc2(x)
    return x