Pseudo-3D Residual Networks算法的pytorch代码

本篇博客是对第三方实现的Pseudo-3D Residual Networks算法的pytorch代码进行介绍，介绍顺序为代码调试顺序，建议先阅读论文或相关博客。
论文：Learning Spatio-Temporal Representation with Pseudo-3D Residual Networks。
代码地址：https://github.com/qijiezhao/pseudo-3d-pytorch

导入必须的模块

from __future__ import print_function
import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from torch.autograd import Variable
import math
from functools import partial

__all__ = ['P3D', 'P3D63', 'P3D131','P3D199']

main函数中主要包括导入模型，准备数据，测试数据，打印结果这几个部分，其中导入模型和准备数据最重要。导入模型部分通过调用P3D199得到199层的P3D网络，另外参数pretrained=true表明导入的是预训练的模型，这样模型的参数就可以通过预训练模型的参数来初始化了，另一个参数num_classes=400是类别数。P3D199后面会详细介绍。
data=torch.autograd.Variable(torch.rand(10,3,16,160,160)).cuda()这一行是随机生成输入数据，第一个维度是10说明该输入数据包含10个clip，其中每个clip包含16帧图像，每帧图像是160*160的3通道图像。最后数据输入模型得到结果out。

if __name__ == '__main__':

    model = P3D199(pretrained=True,num_classes=400)
    model = model.cuda()

# if modality=='Flow', please change the 2nd dimension 3==>2
   data=torch.autograd.Variable(torch.rand(10,3,16,160,160)).cuda()
    out=model(data)
    print (out.size(),out)

调用P3D199函数得到199层的P3D网络。如果调用P3D63或P3D131则得到对应层数的P3D网络，不过由于P3D63和P3D131没有预训练模型，所以在实现中只有调用P3D类导入网络结构这一步。model = P3D(Bottleneck, [3, 8, 36, 3], modality=modality,**kwargs) 这一行通过调用P3D类获得网络结构，该函数的第二个参数[3, 8, 36, 3]表明其结构和ResNet-152对应，于是最后得到的就是199层的P3D网络。因为pretrained设置为True，所以pretrained_file就是准备好的预训练模型的压缩文件，通过weights=torch.load(pretrained_file)['state_dict']这一行读取预训练模型的参数，然后model.load_state_dict(weights)这一行将读取到的预训练模型的参数赋值给model这个网络结构，完成赋值。接下来介绍P3D这个类。

def P3D199(pretrained=False,modality='RGB',**kwargs):
    """construct a P3D199 model based on a ResNet-152-3D model.
    """
    model = P3D(Bottleneck, [3, 8, 36, 3], modality=modality,**kwargs)
    if pretrained==True:
        if modality=='RGB':
            pretrained_file='p3d_rgb_199.checkpoint.pth.tar'
        elif modality=='Flow':
            pretrained_file='p3d_flow_199.checkpoint.pth.tar'
        weights=torch.load(pretrained_file)['state_dict']
        model.load_state_dict(weights)
    return model

P3D这个类是构造网络结构的主体，之前介绍过在PyTorch中定义网络结构的时候都要继承基类torch.nn.Module，这里也是这样。先看看__init__，这是因为前面调用P3D类生成对象的时候会先调用__init__进行初始化。 self.input_channel = 3 if modality=='RGB' else 2这一行表示如果输入是视频帧（也就是图像），那么输入channel就是3，如果输入是optical flow，那么输入channel就是2。self.conv1_custom是网络的第一个卷积层，原来在ResNet中是2D的7*7大小的卷积核。self.maxpool是对最后3维都做了pooling，也就是对这3维都做了减半。self.maxpool_2是对倒数第3维度做了pooling，也就是对该维度做了减半。self.layer1到self.layer4是P3D网络的4个block，这个和ResNet网络中的conv2_x到conv5_x对应，以self.layer1 = self._make_layer(block, 64, layers[0], shortcut_type)为例，通过调用P3D类的make_layer方法来得到该block的层（后面会详细介绍_make_layer方法），一般函数名前面加上表示类内部函数，或者叫私有函数。几个输入的含义如下：block是Bottleneck，layers是一个列表（长度就是block的数量，这里就是4），其中的每个值表示对应的那个block包含多少个重复的residual结构，shortcut_type是residual采取的形式，默认是’B’，指明了ResNet中residual的具体类型。self.avgpool和ResNet网络中最后的7*7的均值池化一样，只不过这里采用的是5*5，主要是因为输入帧的大小是160*160，和ResNet中224*224大小的图像不同。self.layer1和其他3个layer不同的是输入中没有stride=2，这是因为在self.layer1之前已经进行过一次pool操作了，所以这里不需要stride，这和ResNet中每个block的feature map尺寸缩减策略是一致的。for m in self.modules():这个循环是用来对构造好的网络结构进行参数初始化，这里只对卷积层和BN层进行初始化。__init__的最后三行先是指定了网络输入数据的尺寸，这里16表示16帧图像，然后指定了均值和标准差用于数据归一化。接下来介绍P3D类的_make_layer方法。

class P3D(nn.Module):

    def __init__(self, block, layers, modality='RGB',
        shortcut_type='B', num_classes=400,dropout=0.5,ST_struc=('A','B','C')):
        self.inplanes = 64
        super(P3D, self).__init__()
        # self.conv1 = nn.Conv3d(3, 64, kernel_size=7, stride=(1, 2,