Create train.py

train.py

@@ -0,0 +1,76 @@
+import os
+import torch
+import torchvision
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from model import DiffusionModel, UNet
+from torchvision.datasets import CocoCaptions
+import argparse
+from tqdm import tqdm
+
+# Config
+IMAGE_SIZE = 256
+BATCH_SIZE = 16
+EPOCHS = 50
+LR = 2e-5
+TIMESTEPS = 1000
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+def load_coco_dataset():
+    transform = transforms.Compose([
+        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+    ])
+    
+    dataset = CocoCaptions(
+        root='./train2017',
+        annFile='./annotations/captions_train2017.json',
+        transform=transform
+    )
+    
+    dataloader = DataLoader(
+        dataset, 
+        batch_size=BATCH_SIZE, 
+        shuffle=True, 
+        num_workers=4,
+        collate_fn=lambda x: (torch.stack([item[0] for item in x]), [item[1] for item in x])
+    )
+    return dataloader
+
+def train():
+    # Setup
+    model = UNet().to(DEVICE)
+    betas = torch.linspace(1e-4, 0.02, TIMESTEPS).to(DEVICE)
+    diffusion = DiffusionModel(model, betas, DEVICE)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LR)
+    dataloader = load_coco_dataset()
+    
+    # Training loop
+    for epoch in range(EPOCHS):
+        pbar = tqdm(dataloader)
+        for images, captions in pbar:
+            images = images.to(DEVICE)
+            
+            # Flatten captions (5 per image) and repeat images
+            captions = [cap for sublist in captions for cap in sublist]
+            images = images.repeat_interleave(5, dim=0)
+            
+            # Sample timesteps
+            t = torch.randint(0, TIMESTEPS, (images.shape[0],), device=DEVICE).long()
+            
+            # Compute loss
+            loss = diffusion.p_losses(images, captions, t)
+            
+            # Optimize
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            
+            pbar.set_description(f"Epoch {epoch}, Loss: {loss.item():.4f}")
+        
+        # Save checkpoint
+        torch.save(model.state_dict(), f"diffusion_model_epoch_{epoch}.pth")
+
+if __name__ == "__main__":
+    train()