train_model.py

import torch
import torch.nn.functional as F
from torch.optim import Adam
from torch.utils.data import DataLoader

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)

    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