resnet34测试:
import torch
import torch.nn as nn
from torch.autograd import Variable
from torchviz import make_dot
from torchvision import models
dummy_input = Variable(torch.rand(1, 3, 224, 224)) #假设输入13张1*28*28的图片
model = models.resnet34()
y = model(dummy_input)
g = make_dot(y)
g.render('LeNet_model', view=False)
结果生成相应文件