2508181735

.gitattributes

@@ -32,5 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
-data/* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

main.py

@@ -4,7 +4,8 @@ from typing import List, Dict, Any
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torchvision import datasets, transforms
+# from torchvision import datasets -> datasetsライブラリを使うため削除
+from torchvision import transforms
 from torch.utils.data import DataLoader, Subset
 import numpy as np
 import base64
@@ -12,6 +13,7 @@ from io import BytesIO
 from PIL import Image
 import random
 import json
+from datasets import load_dataset # ★★★ Hugging Face datasetsライブラリをインポート
 
 # FastAPIアプリケーションインスタンスを作成
 app = FastAPI()
@@ -22,7 +24,7 @@ class PlayerModel(nn.Module):
         super(PlayerModel, self).__init__()
         self.layers = nn.ModuleList()
         self.architecture_info = []
-        self.hookable_layers = {} # ★★★ nameをキーとしてフック対象レイヤーを保存
+        self.hookable_layers = {}
 
         in_channels = 1
         feature_map_size = 28
@@ -30,10 +32,8 @@ class PlayerModel(nn.Module):
         
         for i, config in enumerate(layer_configs):
             layer_type = config['type']
-            # ユニークな名前を生成
             name = f"{layer_type.lower()}_{len([info for info in self.architecture_info if info['type'] == layer_type])}"
 
-            # 畳み込み/プーリング層
             if layer_type in ['Conv2d', 'MaxPool2d', 'AvgPool2d']:
                 is_flattened = False
                 if layer_type == 'Conv2d':
@@ -44,27 +44,23 @@ class PlayerModel(nn.Module):
                     self.hookable_layers[name] = layer
                     in_channels = out_channels
                     self.architecture_info.append({"type": "Conv2d", "name": name, "shape": [out_channels, feature_map_size, feature_map_size]})
-                else: # MaxPool2d or AvgPool2d
+                else:
                     kernel_size = config['params']['kernel_size']
                     if layer_type == 'MaxPool2d':
                         layer = nn.MaxPool2d(kernel_size=kernel_size, stride=kernel_size)
-                    else: # AvgPool2d
+                    else:
                         layer = nn.AvgPool2d(kernel_size=kernel_size, stride=kernel_size)
                     self.layers.append(layer)
                     self.hookable_layers[name] = layer
                     feature_map_size //= kernel_size
                     self.architecture_info.append({"type": layer_type, "name": name, "shape": [in_channels, feature_map_size, feature_map_size]})
-
-            # 活性化/ドロップアウト層
             elif layer_type in ['ReLU', 'Dropout']:
                 if layer_type == 'ReLU':
                     self.layers.append(nn.ReLU())
-                else: # Dropout
+                else:
                     p = config['params']['p']
                     self.layers.append(nn.Dropout(p=p))
                 self.architecture_info.append({"type": layer_type, "name": name})
-
-            # 構造変更層
             elif layer_type == 'Flatten':
                 if not is_flattened:
                     layer = nn.Flatten()
@@ -74,8 +70,6 @@ class PlayerModel(nn.Module):
                     in_channels = flat_features
                     self.architecture_info.append({"type": "Flatten", "name": name, "shape": [flat_features]})
                     is_flattened = True
-            
-            # 全結合/残差ブロック層
             elif layer_type in ['Linear', 'ResidualBlock']:
                 if not is_flattened:
                     auto_flatten_name = f"auto_flatten_{i}"
@@ -84,28 +78,23 @@ class PlayerModel(nn.Module):
                     in_channels = flat_features
                     self.architecture_info.append({"type": "Flatten", "name": auto_flatten_name, "shape": [flat_features]})
                     is_flattened = True
-                
                 if layer_type == 'Linear':
                     out_features = config['params']['out_features']
                     layer = nn.Linear(in_channels, out_features)
                     in_channels = out_features
-                else: # ResidualBlock
-                    # ★★★ 残差ブロックは次元を維持する線形層として実装
+                else:
                     features = in_channels
                     layer = nn.Linear(features, features)
-                
                 self.layers.append(layer)
                 self.hookable_layers[name] = layer
                 self.architecture_info.append({"type": layer_type, "name": name, "shape": [in_channels]})
 
-        # 最終出力層を強制的に追加
         if not self.layers or not isinstance(self.layers[-1], nn.Linear) or self.layers[-1].out_features != 10:
             if not is_flattened:
                  self.layers.append(nn.Flatten())
                  final_in_features = in_channels * feature_map_size * feature_map_size
             else:
                  final_in_features = in_channels
-            
             output_layer = nn.Linear(final_in_features, 10)
             self.layers.append(output_layer)
             self.hookable_layers["linear_output"] = output_layer
@@ -116,18 +105,45 @@ class PlayerModel(nn.Module):
             x = layer(x)
         return x
 
-# --- グローバル変数 (変更なし) ---
+# --- グローバル変数とデータ準備 (★★★ ここを大幅に変更) ---
 device = torch.device("cpu")
-transform = transforms.ToTensor()
-full_train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
-train_subset = Subset(full_train_dataset, range(1000))
+
+# 1. Hugging Face HubからMNISTデータセットをロード
+mnist_dataset = load_dataset("mnist")
+
+# 2. torchvisionのtransformを定義
+transform = transforms.Compose([
+    transforms.ToTensor(),
+    # transforms.Normalize((0.1307,), (0.3081,)) # 必要に応じて正規化
+])
+
+# 3. データセットにtransformを適用する関数を定義
+def apply_transforms(examples):
+    # PIL Imageのリストをテンソルのリストに変換
+    examples['image'] = [transform(image.convert("L")) for image in examples['image']]
+    return examples
+
+# 4. データセットにtransformを適用
+mnist_dataset.set_transform(apply_transforms)
+
+# 5. DataLoaderを準備
+train_subset = mnist_dataset['train'].select(range(1000))
 train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)
-test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)
-test_images = list(DataLoader(test_dataset, batch_size=1, shuffle=True))
+
+# 6. テスト用の画像リストを作成 (DataLoaderと同じ (image, label) タプルの形式を維持)
+test_images = []
+# メモリ使用量を考慮し、テスト画像は1000個に絞る
+test_subset_for_inference = mnist_dataset['test'].shuffle().select(range(1000))
+for item in test_subset_for_inference:
+    # itemは {'image': tensor, 'label': int} という辞書
+    image_tensor = item['image'].unsqueeze(0)  # バッチ次元 (1) を追加
+    label_tensor = torch.tensor(item['label'])
+    test_images.append((image_tensor, label_tensor))
+
 current_enemy = None
 trained_player_model = None
 
-# --- バックエンドロジック (Eelデコレータを削除) ---
+# --- バックエンドロジック ---
 def get_enemy():
     global current_enemy
     image, label = random.choice(test_images)
@@ -152,8 +168,9 @@ def train_player_model(layer_configs: list):
         
         model.train()
         for epoch in range(3): # 3エポック学習
-            for batch_idx, (data, target) in enumerate(train_loader):
-                data, target = data.to(device), target.to(device)
+            # ★★★ DataLoaderからのデータ受け取り方を変更
+            for batch_idx, batch in enumerate(train_loader):
+                data, target = batch['image'].to(device), batch['label'].to(device)
                 optimizer.zero_grad()
                 output = model(data)
                 loss = loss_fn(output, target)
@@ -223,7 +240,7 @@ def run_inference():
         "weights": weights
     }
 
-# --- FastAPI Endpoints ---
+# --- FastAPI Endpoints (変更なし) ---
 @app.get("/api/get_enemy")
 async def get_enemy_endpoint():
     return get_enemy()
@@ -236,6 +253,5 @@ async def train_player_model_endpoint(layer_configs: List[Dict[str, Any]] = Body
 async def run_inference_endpoint():
     return run_inference()
 
-# --- 静的ファイルの配信 ---
-# フロントエンドのファイル (index.html, style.css, script.js) を配信
+# --- 静的ファイルの配信 (変更なし) ---
 app.mount("/", StaticFiles(directory="web", html=True), name="static")

requirements.txt

@@ -3,4 +3,5 @@ uvicorn[standard]
 torch
 torchvision
 numpy
-Pillow
+Pillow
+datasets