garbage

Navy1228 · May 21, 2020 · a83a572 · a83a572
1 parent 33d32a0
commit a83a572
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/Pytorch-Seg/lesson-4/infer.py b/Pytorch-Seg/lesson-4/infer.py
@@ -27,6 +27,7 @@ def softmax(x):
     fc_inputs = model.fc.in_features
     model.fc = nn.Linear(fc_inputs, 214)
     model = model.cuda()
+    # 加载训练好的模型
     checkpoint = torch.load('model_best_checkpoint_resnet50.pth.tar')
     model.load_state_dict(checkpoint['state_dict'])
     model.eval()
@@ -42,4 +43,4 @@ def softmax(x):
         pred_id = np.argmax(score)
         plt.imshow(src)
         print('预测结果：', labels[pred_id][0])
-        plt.show()
+        plt.show()