1import torch
2from torchvision.models import resnet50
3
4# 加载模型
5model = resnet50(pretrained=False)
6model.fc = torch.nn.Identity()
7model.detection_head = torch.nn.Sequential(
8 torch.nn.Linear(2048, 512),
9 torch.nn.ReLU(),
10 torch.nn.Linear(512, 6) # 4个坐标 + 2个类别
11)
12model.load_state_dict(torch.load('pytorch_model.bin'))
13model.eval()