class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3, stride=1, padding=0),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
)
# 16 input channel image 22x22, 32 output channel image 10x10
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=0),
nn.ReLU(),
nn.MaxPool2d(2)
)
# 分类
self.fc = nn.Sequential(
nn.Linear(32 * 10 * 10, 100)
)
self.classifier = nn.Sequential(
nn.Linear(100, 2)
)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(-1, 32 * 10 * 10)
x = self.fc(x)
x = self.classifier(x)
return x
if __name__ == '__main__':
model = Net()
state_dict = torch.load(r"models\scenejudgment.pth")
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
model.eval()
example = torch.rand(1, 3, 90, 90)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save(r"models\new.pt")
python:3.6.12
pytorch:1.5.0