PyTorch Template

The PyTorch template creates a production-ready ML project with PyTorch, torchvision, and common ML dependencies. The template system provides two variants:

pytorch: Basic inference template for deployment
pytorch-train: Complete training pipeline with advanced features

Quick Start

Create a new PyTorch inference project:

cirron init my-pytorch-model --template pytorch

Create a PyTorch training project:

cirron init my-training-project --template pytorch-train

The model type (classification, regression, etc.) is selected interactively during cirron init.

Project Structure

The PyTorch template generates the following project structure:

my-pytorch-model/
├── src/
│   ├── model.py           # Model architecture definition
│   ├── data_loader.py     # Data loading utilities
│   ├── train.py           # Training script
│   └── inference.py       # Inference/prediction script
├── data/                  # Data directory
├── checkpoints/           # Model checkpoints (created during training)
├── requirements.txt       # Python dependencies
├── Dockerfile            # Container configuration
└── cirron.yaml           # Project configuration

Generated Files

requirements.txt

torch>=2.0.0
torchvision>=0.15.0
numpy>=1.21.0
scikit-learn>=1.3.0
matplotlib>=3.5.0
tqdm>=4.64.0
Pillow>=9.0.0
requests>=2.28.0

src/model.py

The model architecture is automatically generated based on your model type:

Classification Model

import torch
import torch.nn as nn

class ClassificationModel(nn.Module):
    def __init__(self, num_classes=10):
        super(ClassificationModel, self).__init__()
        self.num_classes = num_classes
        
        # CNN layers
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        
        # Pooling and dropout
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout(0.25)
        
        # Fully connected layers
        self.fc1 = nn.Linear(128 * 28 * 28, 512)
        self.fc2 = nn.Linear(512, num_classes)
        
        # Activation
        self.relu = nn.ReLU()
    
    def forward(self, x):
        # Convolutional layers
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = self.pool(self.relu(self.conv3(x)))
        
        # Flatten
        x = x.view(x.size(0), -1)
        
        # Fully connected layers
        x = self.dropout(self.relu(self.fc1(x)))
        x = self.fc2(x)
        
        return x

Regression Model

import torch
import torch.nn as nn

class RegressionModel(nn.Module):
    def __init__(self, input_size=784):
        super(RegressionModel, self).__init__()
        
        # Fully connected layers
        self.fc1 = nn.Linear(input_size, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 128)
        self.fc4 = nn.Linear(128, 1)
        
        # Activation and dropout
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        # Flatten input if needed
        if len(x.shape) > 2:
            x = x.view(x.size(0), -1)
        
        # Fully connected layers
        x = self.dropout(self.relu(self.fc1(x)))
        x = self.dropout(self.relu(self.fc2(x)))
        x = self.dropout(self.relu(self.fc3(x)))
        x = self.fc4(x)
        
        return x

src/data_loader.py

Data loading utilities for PyTorch:

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import os
from PIL import Image
import numpy as np

class CustomDataset(Dataset):
    def __init__(self, data_path, transform=None):
        self.data_path = data_path
        self.transform = transform
        self.samples = self._load_samples()
    
    def _load_samples(self):
        # TODO: Implement based on your data structure
        # This is a placeholder - customize for your data
        samples = []
        # Add your data loading logic here
        return samples
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        # TODO: Implement based on your data structure
        # This is a placeholder - customize for your data
        sample = self.samples[idx]
        
        if self.transform:
            sample = self.transform(sample)
        
        return sample

def get_data_loaders(config):
    """Get training and validation data loaders"""
    
    # Define transforms
    if config['model_type'] == 'classification':
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                               std=[0.229, 0.224, 0.225])
        ])
    else:
        transform = transforms.Compose([
            transforms.ToTensor()
        ])
    
    # Create datasets
    train_dataset = CustomDataset(
        os.path.join(config['data_path'], 'train'),
        transform=transform
    )
    
    val_dataset = CustomDataset(
        os.path.join(config['data_path'], 'val'),
        transform=transform
    )
    
    # Create data loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=config['batch_size'],
        shuffle=True,
        num_workers=4
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=config['batch_size'],
        shuffle=False,
        num_workers=4
    )
    
    return train_loader, val_loader

src/train.py

Complete training script with best practices:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import os
from tqdm import tqdm

from model import ClassificationModel  # or RegressionModel
from data_loader import get_data_loaders

class Trainer:
    def __init__(self, config):
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Initialize model
        self.model = ClassificationModel()  # or RegressionModel()
        self.model.to(self.device)
        
        # Loss and optimizer
        self.criterion = self.get_criterion()
        self.optimizer = optim.Adam(self.model.parameters(), lr=config['learning_rate'])
        
        # Data loaders
        self.train_loader, self.val_loader = get_data_loaders(config)
        
        # Metrics
        self.train_losses = []
        self.val_losses = []
    
    def get_criterion(self):
        """Get appropriate loss function"""
        if self.config['model_type'] == 'classification':
            return nn.CrossEntropyLoss()
        elif self.config['model_type'] == 'regression':
            return nn.MSELoss()
        else:
            return nn.MSELoss()  # Default
    
    def train_epoch(self):
        """Train for one epoch"""
        self.model.train()
        total_loss = 0
        
        for batch_idx, (data, target) in enumerate(tqdm(self.train_loader, desc="Training")):
            data, target = data.to(self.device), target.to(self.device)
            
            self.optimizer.zero_grad()
            output = self.model(data)
            loss = self.criterion(output, target)
            loss.backward()
            self.optimizer.step()
            
            total_loss += loss.item()
        
        return total_loss / len(self.train_loader)
    
    def validate(self):
        """Validate the model"""
        self.model.eval()
        total_loss = 0
        correct = 0
        
        with torch.no_grad():
            for data, target in tqdm(self.val_loader, desc="Validation"):
                data, target = data.to(self.device), target.to(self.device)
                output = self.model(data)
                loss = self.criterion(output, target)
                total_loss += loss.item()
                
                # Calculate accuracy for classification
                if self.config['model_type'] == 'classification':
                    pred = output.argmax(dim=1, keepdim=True)
                    correct += pred.eq(target.view_as(pred)).sum().item()
        
        avg_loss = total_loss / len(self.val_loader)
        accuracy = correct / len(self.val_loader.dataset) if self.config['model_type'] == 'classification' else None
        
        return avg_loss, accuracy
    
    def save_checkpoint(self, epoch, val_loss, is_best=False):
        """Save model checkpoint"""
        os.makedirs('checkpoints', exist_ok=True)
        
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'val_loss': val_loss,
            'config': self.config
        }
        
        checkpoint_path = f'checkpoints/checkpoint_epoch_{epoch}.pth'
        torch.save(checkpoint, checkpoint_path)
        
        if is_best:
            torch.save(checkpoint, 'checkpoints/best_model.pth')
    
    def train(self):
        """Main training loop"""
        best_val_loss = float('inf')
        
        for epoch in range(self.config['num_epochs']):
            print(f"\nEpoch {epoch+1}/{self.config['num_epochs']}")
            
            # Train
            train_loss = self.train_epoch()
            self.train_losses.append(train_loss)
            
            # Validate
            val_loss, accuracy = self.validate()
            self.val_losses.append(val_loss)
            
            # Print metrics
            print(f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")
            if accuracy is not None:
                print(f"Val Accuracy: {accuracy:.4f}")
            
            # Save checkpoint
            is_best = val_loss < best_val_loss
            if is_best:
                best_val_loss = val_loss
            
            self.save_checkpoint(epoch, val_loss, is_best)

if __name__ == "__main__":
    config = {
        'batch_size': 32,
        'learning_rate': 0.001,
        'num_epochs': 10,
        'model_type': 'classification',  # or 'regression'
        'data_path': 'data/',
    }
    
    trainer = Trainer(config)
    trainer.train()

src/inference.py

Production-ready inference script:

import torch
import torch.nn.functional as F
from PIL import Image
import numpy as np
from model import ClassificationModel  # or RegressionModel

class ModelInference:
    def __init__(self, model_path: str = None):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = ClassificationModel()  # or RegressionModel()
        
        if model_path:
            self.load_model(model_path)
        
        self.model.to(self.device)
        self.model.eval()
    
    def load_model(self, model_path: str):
        """Load trained model weights"""
        checkpoint = torch.load(model_path, map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        print(f"Model loaded from {model_path}")
    
    def preprocess(self, input_data):
        """Preprocess input data"""
        # TODO: Implement preprocessing based on your model type
        if isinstance(input_data, Image.Image):
            # For image inputs
            input_data = input_data.resize((224, 224))
            input_tensor = torch.tensor(np.array(input_data)).float()
            input_tensor = input_tensor.permute(2, 0, 1).unsqueeze(0)
        else:
            # For other data types
            input_tensor = torch.tensor(input_data).float()
            if len(input_tensor.shape) == 1:
                input_tensor = input_tensor.unsqueeze(0)
        
        return input_tensor.to(self.device)
    
    def predict(self, input_data):
        """Make prediction"""
        with torch.no_grad():
            input_tensor = self.preprocess(input_data)
            output = self.model(input_tensor)
            
            # Apply appropriate activation based on model type
            if hasattr(self.model, 'num_classes') and self.model.num_classes > 1:
                predictions = F.softmax(output, dim=1)
            else:
                predictions = output
            
            return predictions.cpu().numpy()

if __name__ == "__main__":
    # Example usage
    inference = ModelInference()
    
    # TODO: Replace with actual input
    sample_input = torch.randn(1, 3, 224, 224)  # Example for image input
    result = inference.predict(sample_input)
    print(f"Prediction: {result}")

Model Types

The PyTorch template supports different model types:

Classification

Use case: Image classification, text classification, multi-class problems
Output: Class probabilities
Loss function: CrossEntropyLoss
Activation: Softmax

Regression

Use case: Price prediction, time series forecasting, continuous value prediction
Output: Continuous values
Loss function: MSELoss
Activation: None (linear output)

Training Configuration

Default training configuration:

config = {
    'batch_size': 32,
    'learning_rate': 0.001,
    'num_epochs': 10,
    'model_type': 'classification',  # or 'regression'
    'data_path': 'data/',
}

Usage Examples

Basic Training

# Navigate to your project
cd my-pytorch-model

# Train the model
python src/train.py

Custom Training

# Modify config in src/train.py
config = {
    'batch_size': 64,
    'learning_rate': 0.0001,
    'num_epochs': 50,
    'model_type': 'classification',
    'data_path': 'data/',
}

Inference

from src.inference import ModelInference

# Load trained model
inference = ModelInference('checkpoints/best_model.pth')

# Make prediction
result = inference.predict(your_input_data)
print(f"Prediction: {result}")

Customization

Adding Custom Models

Modify src/model.py to add your custom architecture
Update the model import in src/train.py and src/inference.py
Adjust the data preprocessing in src/data_loader.py

Custom Data Loading

Update the CustomDataset class in src/data_loader.py
Implement the _load_samples() and __getitem__() methods
Adjust transforms based on your data format

Hyperparameter Tuning

Modify the config dictionary in src/train.py:

config = {
    'batch_size': 64,           # Larger batches for more stable training
    'learning_rate': 0.0001,    # Lower learning rate for fine-tuning
    'num_epochs': 100,          # More epochs for better convergence
    'model_type': 'classification',
    'data_path': 'data/',
}

Best Practices

Data Organization

data/
├── train/
│   ├── class1/
│   ├── class2/
│   └── ...
└── val/
    ├── class1/
    ├── class2/
    └── ...

Model Checkpoints

Checkpoints are saved in the checkpoints/ directory
Best model is saved as checkpoints/best_model.pth
Each epoch checkpoint includes model state, optimizer state, and configuration

GPU Support

Automatically detects CUDA availability
Falls back to CPU if GPU is not available
Use torch.cuda.is_available() to check GPU status

Next Steps

TensorFlow Template - Alternative deep learning framework
scikit-learn Template - Traditional ML algorithms
Training Guide - Advanced training techniques
Deployment Guide - Deploy your trained model

Cirron CLI

Commands

Utils

Templates

PyTorch

PyTorch Template

Quick Start

Project Structure

Generated Files

requirements.txt

src/model.py

Classification Model

Regression Model

src/data_loader.py

src/train.py

src/inference.py

Model Types

Classification

Regression

Training Configuration

Usage Examples

Basic Training

Custom Training

Inference

Customization

Adding Custom Models

Custom Data Loading

Hyperparameter Tuning

Best Practices

Data Organization

Model Checkpoints

GPU Support

Next Steps

Cirron CLI

Commands

Utils

Templates

Documentation Index

​PyTorch Template

​Quick Start

​Project Structure

​Generated Files

​requirements.txt

​src/model.py

​Classification Model

​Regression Model

​src/data_loader.py

​src/train.py

​src/inference.py

​Model Types

​Classification

​Regression

​Training Configuration

​Usage Examples

​Basic Training

​Custom Training

​Inference

​Customization

​Adding Custom Models

​Custom Data Loading

​Hyperparameter Tuning

​Best Practices

​Data Organization

​Model Checkpoints

​GPU Support

​Next Steps

PyTorch Template

Quick Start

Project Structure

Generated Files

requirements.txt

src/model.py

Classification Model

Regression Model

src/data_loader.py

src/train.py

src/inference.py

Model Types

Classification

Regression

Training Configuration

Usage Examples

Basic Training

Custom Training

Inference

Customization

Adding Custom Models

Custom Data Loading

Hyperparameter Tuning

Best Practices

Data Organization

Model Checkpoints

GPU Support

Next Steps