如何知道Pytorch模型的输入/输出层名称和大小？

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)

input_name = list(model._modules.keys())[0]  # 获取模型的第一个子模块的名称
input_size = model._modules[input_name].weight.size()  # 获取输入层的大小

output_name = list(model._modules.keys())[-1]  # 获取模型的最后一个子模块的名称
output_size = model._modules[output_name].weight.size()  # 获取输出层的大小