finish exp2

exp2/modified/predict.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+import sys
+import torch
+from train import CNN
+
+sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
+import gpu_utils
+
+
+class PredictResult:
+
+    possibilities: torch.Tensor
+
+    def __init__(self, possibilities: torch.Tensor):
+        self.possibilities = possibilities
+
+    def chosen_number(self) -> int:
+        """获取最终选定的数字"""
+        # 依然是找最大的那个index
+        _, prediction = self.possibilities.max(1)
+        return prediction.item()
+
+    def number_possibilities(self) -> list[float]:
+        """获取每个数字出现的概率"""
+        return list(self.possibilities[0][i].item() for i in range(10))
+
+class Predictor:
+    device: torch.device
+    cnn: CNN
+
+    def __init__(self):
+        self.device = gpu_utils.get_gpu_device()
+        self.cnn = CNN().to(self.device)
+
+        # 加载保存好的模型参数
+        file_path = Path(__file__).resolve().parent.parent / 'models' / 'cnn.pth'
+        self.cnn.load_state_dict(torch.load(file_path))
+
+    def predict(self, image: list[list[bool]]) -> PredictResult:
+        input = torch.Tensor(image).float().to(self.device)
+        assert(input.dim() == 2)
+        assert(input.size(0) == 28)
+        assert(input.size(1) == 28)
+
+        # 为了满足要求，要在第一维度挤出2下
+        # 一次是灰度通道，一次是批次。
+        # 相当于batch size = 1的计算
+        input = input.unsqueeze(0).unsqueeze(0)
+       
+        # 预测
+        with torch.no_grad():
+            output = self.cnn(input)
+            return PredictResult(output)
+