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里实现的modelclassResNet50(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)defforward(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 mapifnot 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()