Update train.py

Author: Snow0821

Experiments/rl_nand_network/train.py

@@ -5,7 +5,6 @@
 from torchvision import datasets, transforms
 from torch.utils.data import DataLoader
 
-# 1️⃣ 데이터 전처리 및 로더
 transform = transforms.Compose([
     transforms.ToTensor(),
     transforms.Lambda(lambda x: (x > 0).float())  # Binarize
@@ -14,19 +13,18 @@
 train_dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
 train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
 
-# 2️⃣ SparseBtnn_Selector 정의
 def check_binary_tensor(tensor, tensor_name="tensor"):
     assert torch.all((tensor == 0) | (tensor == 1)), f"{tensor_name} contains values other than -1 or 1! as {tensor}"
 
 def my_hard_mask(logits):
     softmax_probs = torch.softmax(logits, dim=1)
     hard_mask = torch.zeros_like(softmax_probs).scatter_(1, softmax_probs.argmax(dim=1, keepdim=True), 1.0)
-    return hard_mask - softmax_probs.detach() + softmax_probs  # STE 적용
+    return hard_mask - softmax_probs.detach() + softmax_probs  # STE
 
 def my_action_mask(logits, action):
     softmax_probs = torch.softmax(logits, dim=1)
     hard_mask = torch.zeros_like(softmax_probs).scatter_(1, action, 1.0)
-    return hard_mask - softmax_probs.detach() + softmax_probs  # STE 적용
+    return hard_mask - softmax_probs.detach() + softmax_probs  # STE
 
 class SparseBtnn_Selector(nn.Module):
     def __init__(self, x, y):
@@ -41,18 +39,14 @@ def forward(self, x, reward=None, train_mode=False):
         if train_mode and reward is not None:
             action = torch.multinomial(action_probs, num_samples=1)
             policy_loss = -10 * torch.log(action_probs.gather(1, action)) * reward
-            policy_loss.mean().backward(retain_graph=True)  # retain_graph 사용
+            policy_loss.mean().backward(retain_graph=True)  # retain_graph
             mask = my_action_mask(selector_logits, action)
         else:
-            mask = my_hard_mask(selector_logits)  # STE 적용 마스크
+            mask = my_hard_mask(selector_logits)  # had mask
         out = F.linear(x, mask)
         check_binary_tensor(out, "selector")
         return out
 
-
-
-
-# 3️⃣ Compacted_Nand 정의
 class Compacted_Nand(nn.Module):
     def __init__(self, x, y):
         super().__init__()
@@ -66,12 +60,9 @@ def forward(self, x, reward=None, train_mode=False):
         check_binary_tensor(out, "nand")
         return out
 
-# 4️⃣ 보상 함수 정의
 def get_reward(prediction, target):
-    """정확한 예측에는 +1 보상, 틀린 예측에는 -1 처벌"""
     return torch.where(prediction == target, torch.tensor(1.0, device=prediction.device), torch.tensor(-1.0, device=prediction.device))
 
-# 5️⃣ 학습 루프
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = Compacted_Nand(28*28, 10).to(device)
 optimizer = optim.Adam(model.parameters(), lr=0.001)
@@ -85,13 +76,8 @@ def get_reward(prediction, target):
     for batch_idx, (data, target) in enumerate(train_loader):
         data, target = data.to(device).view(data.size(0), -1), target.to(device)
 
-        # 1️⃣ 순전파 (logits 생성)
         logits = model(data, train_mode=True)
-        
-        # 2️⃣ 보상 계산 (logits가 정의된 후)
         reward = get_reward(logits.argmax(dim=1), target)
-        
-        # 3️⃣ 손실 계산 및 역전파
         loss = criterion(logits, target)
         optimizer.zero_grad()
         loss.backward()