import models from torchsummary import summary model = models.resnet50(pretrained=False, num_classes=7, scale=1) print(model) summary(model, (3, 640, 640))
使用torchsummay模块,pip install torchsummary
summayr使用方法,加入模型,以及模型输入tensor的形状即可,结尾会给出模型的权重大小,反向传播所需内存大小
使用Print同样可以直接打印模型结构,两者各有优劣。