Upload train_mnist.py with huggingface_hub

train_mnist.py

@@ -0,0 +1,166 @@
+# Training Links:
+# W&B Run: https://wandb.ai/ivanleo97-freelance/mnist-modal/runs/tu4yqtvi
+# Hugging Face Model: https://huggingface.co/ivanleomk/mnist-modal
+
+import modal
+
+app = modal.App("mnist-training")
+
+# We use an image with torch, torchvision, wandb and huggingface_hub installed.
+# We will import these inline inside the function to respect the user's request
+# and avoid needing them installed locally.
+image = modal.Image.debian_slim().pip_install(
+    "torch", "torchvision", "wandb", "huggingface_hub"
+)
+
+@app.function(
+    image=image, 
+    gpu="A100", 
+    timeout=3600, 
+    secrets=[
+        modal.Secret.from_dict({"WANDB_API_KEY": "YOUR_WANDB_API_KEY"}),
+        modal.Secret.from_dict({"HF_TOKEN": "YOUR_HF_TOKEN"})
+    ]
+)
+def train():
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    import torch.optim as optim
+    from torchvision import datasets, transforms
+    import wandb
+    import os
+    from huggingface_hub import HfApi
+    
+    class Net(nn.Module):
+        def __init__(self):
+            super(Net, self).__init__()
+            self.conv1 = nn.Conv2d(1, 32, 3, 1)
+            self.conv2 = nn.Conv2d(32, 64, 3, 1)
+            self.dropout1 = nn.Dropout(0.25)
+            self.dropout2 = nn.Dropout(0.5)
+            self.fc1 = nn.Linear(9216, 128)
+            self.fc2 = nn.Linear(128, 10)
+
+        def forward(self, x):
+            x = F.relu(self.conv1(x))
+            x = F.relu(self.conv2(x))
+            x = F.max_pool2d(x, 2)
+            x = self.dropout1(x)
+            x = torch.flatten(x, 1)
+            x = F.relu(self.fc1(x))
+            x = self.dropout2(x)
+            x = self.fc2(x)
+            return F.log_softmax(x, dim=1)
+            
+    # Initialize W&B
+    wandb.init(project="mnist-modal", config={
+        "learning_rate": 1.0,
+        "epochs": 5,
+        "batch_size": 64
+    })
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
+    model = Net().to(device)
+    # Using Adadelta as in standard pytorch mnist example
+    optimizer = optim.Adadelta(model.parameters(), lr=wandb.config.learning_rate)
+    
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.1307,), (0.3081,))
+    ])
+    
+    print("Downloading dataset...")
+    train_loader = torch.utils.data.DataLoader(
+        datasets.MNIST('/tmp/data', train=True, download=True, transform=transform),
+        batch_size=wandb.config.batch_size, shuffle=True)
+        
+    test_loader = torch.utils.data.DataLoader(
+        datasets.MNIST('/tmp/data', train=False, transform=transform),
+        batch_size=1000, shuffle=True)
+        
+    print("Starting training...")
+    for epoch in range(1, wandb.config.epochs + 1):
+        model.train()
+        train_loss = 0
+        for batch_idx, (data, target) in enumerate(train_loader):
+            data, target = data.to(device), target.to(device)
+            optimizer.zero_grad()
+            output = model(data)
+            loss = F.nll_loss(output, target)
+            loss.backward()
+            optimizer.step()
+            train_loss += loss.item()
+            
+            if batch_idx % 100 == 0:
+                print(f"Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} "
+                      f"({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}")
+                
+        train_loss /= len(train_loader)
+        
+        # Test
+        model.eval()
+        test_loss = 0
+        correct = 0
+        with torch.no_grad():
+            for data, target in test_loader:
+                data, target = data.to(device), target.to(device)
+                output = model(data)
+                test_loss += F.nll_loss(output, target, reduction='sum').item()
+                pred = output.argmax(dim=1, keepdim=True)
+                correct += pred.eq(target.view_as(pred)).sum().item()
+                
+        test_loss /= len(test_loader.dataset)
+        accuracy = 100. * correct / len(test_loader.dataset)
+        
+        print(f"\nEpoch {epoch} summary: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n")
+        
+        wandb.log({
+            "epoch": epoch,
+            "train_loss": train_loss,
+            "test_loss": test_loss,
+            "accuracy": accuracy
+        })
+        
+    print(f"Final test accuracy: {accuracy:.2f}%")
+    
+    # Save model
+    model_path = "/tmp/mnist_model.pth"
+    torch.save(model.state_dict(), model_path)
+    print(f"Model saved to {model_path}")
+    
+    # Upload to HF
+    try:
+        api = HfApi()
+        user_info = api.whoami(token=os.environ["HF_TOKEN"])
+        username = user_info["name"]
+        repo_id = f"{username}/mnist-modal"
+        
+        print(f"Uploading model to Hugging Face repo: {repo_id}")
+        
+        try:
+            api.create_repo(repo_id=repo_id, repo_type="model", token=os.environ["HF_TOKEN"], exist_ok=True)
+        except Exception as e:
+            print(f"Repo might already exist or error: {e}")
+            
+        api.upload_file(
+            path_or_fileobj=model_path,
+            path_in_repo="mnist_model.pth",
+            repo_id=repo_id,
+            repo_type="model",
+            token=os.environ["HF_TOKEN"]
+        )
+        print("Model uploaded successfully!")
+    except Exception as e:
+        print(f"Failed to upload to HF: {e}")
+        
+    wandb.finish()
+    
+    return {"accuracy": accuracy}
+
+@app.local_entrypoint()
+def main():
+    train.remote()
+