Update train_model.py

train_model.py

@@ -1,16 +1,29 @@
-for epoch in range(num_epochs):
-    optimizer.zero_grad()
-    anchor_vec = product_model(anchor_data)
-    positive_vec = product_model(positive_data)
-    negative_vec = product_model(negative_data)
+import torch
+import torch.nn.functional as F
+from torch.optim import Adam
+from torch.utils.data import DataLoader
 
-    # 트립렛 손실 계산
-    positive_distance = F.pairwise_distance(anchor_vec, positive_vec)
-    negative_distance = F.pairwise_distance(anchor_vec, negative_vec)
-    triplet_loss = torch.clamp(positive_distance - negative_distance + margin, min=0).mean()
-    
-    # 역전파와 최적화
-    triplet_loss.backward()
-    optimizer.step()
+def train_triplet_model(product_model, anchor_data, positive_data, negative_data, num_epochs=10, learning_rate=0.001, margin=1.0):
+    optimizer = Adam(product_model.parameters(), lr=learning_rate)
 
-    print(f"Epoch {epoch + 1}, Loss: {triplet_loss.item()}")
+    for epoch in range(num_epochs):
+        product_model.train()
+        optimizer.zero_grad()
+
+        # Forward pass
+        anchor_vec = product_model(anchor_data)
+        positive_vec = product_model(positive_data)
+        negative_vec = product_model(negative_data)
+
+        # Triplet loss calculation
+        positive_distance = F.pairwise_distance(anchor_vec, positive_vec)
+        negative_distance = F.pairwise_distance(anchor_vec, negative_vec)
+        triplet_loss = torch.clamp(positive_distance - negative_distance + margin, min=0).mean()
+
+        # Backward pass and optimization
+        triplet_loss.backward()
+        optimizer.step()
+
+        print(f"Epoch {epoch + 1}, Loss: {triplet_loss.item()}")
+
+    return product_model