加载:
import torch
import torchvision.models as models
from torchvision.models import ResNet18_Weights
# 加载预训练的 ResNet-18 模型,使用新的 weights 参数
model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
# 设置模型为评估模式
model.eval()
# 生成一个输入张量
input_image = torch.randn(1, 3, 224, 224)
# 进行推理
with torch.no_grad():
output = model(input_image)
print(output.shape) # 输出形状为 (1, 1000)
运行输出:
szj@bcc:~/model$ python3 load_resnet18.py
torch.Size([1, 1000])