利用tensorwatch对pytorch网络可视化及统计参数

代码

import sys
import torch
import tensorwatch as tw
import torchvision.models

alexnet_model = torchvision.models.alexnet()
tw.draw_model(alexnet_model, [1, 3, 224, 224])
tw.model_stats(alexnet_model, [1, 3, 224, 224])

结果