基于表征学习的ReID代码实践（4）

model.py：提取特征的模型（ResNet50）

from __future__ import absolute_import

import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from IPython import embed

'''
任何一个神经网络的定义都要继承自nn.Module
一个神经网络的定义包括两个部分
__init__（定义网络结构）和forward（定义前传）
'''


# 方便起见，直接使用resnet时我们可以调用torch里实现的model
class ResNet50(nn.Module):
    def __init__(self, num_classes, **kwargs):
        super(ResNet50, self).__init__()
        resnet50 = torchvision.models.resnet50(pretrained=True)

        # 把resnet50的最后两层（pooling和fc）去掉
        self.base = nn.Sequential(*list(resnet50.children())[:-2])
        self.classifier = nn.Linear(2048, num_classes)

    def forward(self, x):
        x = self.base(x)
        # 得到32*2048*1*1
        x = F.avg_pool2d(x,x.size()[2:])
        # 得到32*2048
        f = x.view(x.size(0),-1)

        # 对特征进行归一化
        #f = 1.*f/(torch.norm(f,2,dim=-1,keepdim=True).expend_as(f) + 1e-12)

        # 如果不是训练，则只输出feature map
        if not self.training:
            return f
        y = self.classifier(f)
        return y



if __name__ == '__main__':
    model = ResNet50(num_classes=751)
    imgs = torch.Tensor(32, 3, 256, 128)
    f = model(imgs)
    embed()