rule

AI/ML Python Projects — Cursor Rules

Production ML: PyTorch, data pipelines, experiment tracking, and deployment

Project Context

You are building an AI/ML project with Python. The project uses PyTorch for model training, handles data pipelines with proper validation, tracks experiments systematically, and follows production ML engineering practices. Code is type-hinted, tested, and reproducible.

Project Structure

project/
  src/
    data/
      datasets.py           # PyTorch Dataset classes
      transforms.py         # Data augmentation and preprocessing
      loaders.py            # DataLoader configurations
      validation.py         # Data quality checks
    models/
      architectures/        # Model definitions
        resnet.py
        transformer.py
      losses.py             # Custom loss functions
      metrics.py            # Evaluation metrics
    training/
      trainer.py            # Training loop
      callbacks.py          # Training callbacks (early stopping, checkpointing)
      optimizers.py         # Optimizer configurations
    inference/
      predictor.py          # Inference pipeline
      postprocess.py        # Output postprocessing
    utils/
      config.py             # Configuration management
      logging.py            # Experiment logging
      reproducibility.py    # Seed setting, deterministic mode
  configs/
    train_config.yaml       # Training hyperparameters
    model_config.yaml       # Model architecture config
  scripts/
    train.py                # Training entry point
    evaluate.py             # Evaluation script
    export.py               # Model export (ONNX, TorchScript)
  notebooks/
    exploration.ipynb       # Data exploration (not production code)
  tests/
    test_data.py
    test_models.py
    test_training.py

PyTorch Model Patterns

• Inherit from nn.Module. Always call super().__init__().
• Type hint all method signatures:

  class ClassificationHead(nn.Module):
      def __init__(self, in_features: int, num_classes: int, dropout: float = 0.1) -> None:
          super().__init__()
          self.dropout = nn.Dropout(dropout)
          self.fc = nn.Linear(in_features, num_classes)
  
      def forward(self, x: torch.Tensor) -> torch.Tensor:
          x = self.dropout(x)
          return self.fc(x)
  

• Use nn.Sequential or nn.ModuleList for dynamic layer construction — never regular Python lists.
• Register buffers for non-parameter tensors: self.register_buffer('mean', torch.zeros(3)).
• Use @torch.no_grad() for inference methods.
• DON'T: Use numpy operations inside forward() — they break autograd.
• DON'T: Create tensors in forward() without sending to the correct device.