faster

app.py

@@ -649,7 +649,7 @@ def start_training_simple(dataset_size: str, res_epochs: int, vit_epochs: int):
             log_message += f"\n🚀 Starting ResNet training on {dataset_size} samples...\n"
             resnet_result =             subprocess.run([
                 "python", "train_resnet.py", "--data_root", DATASET_ROOT, "--epochs", str(res_epochs),
-                "--batch_size", "20", "--out", os.path.join(export_dir, "resnet_item_embedder.pth")
+                "--batch_size", "8", "--out", os.path.join(export_dir, "resnet_item_embedder.pth")
             ] + dataset_args, capture_output=True, text=True, check=False)
             
             if resnet_result.returncode == 0:
@@ -675,7 +675,7 @@ def start_training_simple(dataset_size: str, res_epochs: int, vit_epochs: int):
             log_message += f"\n🚀 Starting ViT training on {dataset_size} samples...\n"
             vit_result = subprocess.run([
                 "python", "train_vit_triplet.py", "--data_root", DATASET_ROOT, "--epochs", str(vit_epochs),
-                "--batch_size", "20", "--export", os.path.join(export_dir, "vit_outfit_model.pth")
+                "--batch_size", "8", "--export", os.path.join(export_dir, "vit_outfit_model.pth")
             ] + dataset_args, capture_output=True, text=True, check=False)
             
             if vit_result.returncode == 0:
@@ -827,7 +827,7 @@ with gr.Blocks(fill_height=True, title="Dressify - Advanced Outfit Recommendatio
                 
                 # Training parameters
                 resnet_epochs = gr.Slider(1, 100, value=20, step=1, label="Epochs")
-                resnet_batch_size = gr.Slider(8, 128, value=20, step=8, label="Batch Size")
+                resnet_batch_size = gr.Slider(8, 128, value=8, step=8, label="Batch Size")
                 resnet_lr = gr.Slider(1e-5, 1e-2, value=1e-3, step=1e-5, label="Learning Rate")
                 resnet_optimizer = gr.Dropdown(
                     choices=["adamw", "adam", "sgd", "rmsprop"], 
@@ -849,7 +849,7 @@ with gr.Blocks(fill_height=True, title="Dressify - Advanced Outfit Recommendatio
                 
                 # Training parameters
                 vit_epochs = gr.Slider(1, 100, value=30, step=1, label="Epochs")
-                vit_batch_size = gr.Slider(4, 64, value=20, step=4, label="Batch Size")
+                vit_batch_size = gr.Slider(4, 64, value=8, step=4, label="Batch Size")
                 vit_lr = gr.Slider(1e-5, 1e-2, value=5e-4, step=1e-5, label="Learning Rate")
                 vit_optimizer = gr.Dropdown(
                     choices=["adamw", "adam", "sgd", "rmsprop"], 

train_resnet.py

@@ -21,7 +21,7 @@ def parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser()
     p.add_argument("--data_root", type=str, default=os.getenv("POLYVORE_ROOT", "/home/user/app/data/Polyvore"))
     p.add_argument("--epochs", type=int, default=20)
-    p.add_argument("--batch_size", type=int, default=20)
+    p.add_argument("--batch_size", type=int, default=8)
     p.add_argument("--lr", type=float, default=1e-3)
     p.add_argument("--embedding_dim", type=int, default=512)
     p.add_argument("--out", type=str, default="models/exports/resnet_item_embedder.pth")

train_vit_triplet.py

@@ -22,7 +22,7 @@ def parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser()
     p.add_argument("--data_root", type=str, default=os.getenv("POLYVORE_ROOT", "/home/user/app/data/Polyvore"))
     p.add_argument("--epochs", type=int, default=30)
-    p.add_argument("--batch_size", type=int, default=20)
+    p.add_argument("--batch_size", type=int, default=8)
     p.add_argument("--lr", type=float, default=5e-4)
     p.add_argument("--embedding_dim", type=int, default=512)
     p.add_argument("--triplet_margin", type=float, default=0.3)