1. 方法一:print_network(net)函数,输出任意模型网络结构
用print_network(net)函数,输出网络结构。代码示例如下。class类里的模型换成任意想要print的结构就行了。
import torch
import torch.nn as nn
import torch.nn.functional as F
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
def print_network(net):
num_params = 0
for param in net.parameters():
num_params += param.numel()
print(net)
print('Total number of parameters: %d' % num_params)
class ModelName(nn.Module):
'''在这里附上要print的模型代码'''
model_test = ModelName()
model_test.to(device)
print('---------- Networks initialized -------------')
print_network(model_test)
print('-----------------------------------------------')
2. 方法二:输出torch官方提供的VGG模型结构和参数
用torch官方提供的已经训练好的模型,只需要从torchvision模块导入:
import torchvision.models as models
import torch
model = models.vgg16(pretrained=True)
feature = torch.nn.Sequential(*list(model.children())[:])
print(feature)