输出网络结构

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)