torchvision 加载模型/查看计算参数/更改输出层

import torch
from torchvision.models import resnet18,resnet34,resnet50,vgg11_bn,inception_v3
from thop import profile
from torchsummary import summary

#加载并查看参数与运算复杂度
model = vgg11_bn(pretrained=False).cuda() #加载模型，未预训练
input = torch.randn(1, 3, 128, 128).cuda()
flops, params = profile(model, inputs=(input, )) #计算参数与计算复杂度
flops, params = clever_format([flops, params], "%.8f")  #按格式输出
print(flops, params)


#查看模型具体结构
print(model) #直接打印模型
summary(model,input_size=(3,128,128)) #格式化打印，input_size = (channel,H,W)


#更改输出层，默认输出为[1000,],通过上面打印的信息查看模型最后一层的名字，可能是fc/classifier[-1]
#model.fc = torch.nn.Linear(model.fc.in_features,num_class).cuda() 
#model.fc.in_features是最后一层的输入大小,num_class为希望的输出大小
model.classifier[-1] = torch.nn.Linear(4096,2).cuda() #以VGG 为例
summary(model,input_size=(3,128,128)) #查看更改后的模型