🌟 We are Open Source! Check out our repository on GitHub

PyTorch Project Structure

A scalable and maintainable folder structure for PyTorch applications

File Structure

  • src
    • train.py
    • test.py
    • inference.py
    • export.py
    • models
      • __init__.py
      • resnet.py
      • transformer.py
      • base_model.py
      • layers.py
    • data
      • __init__.py
      • dataset.py
      • transforms.py
      • dataloader.py
      • augmentation.py
    • training
      • __init__.py
      • trainer.py
      • optimizer.py
      • scheduler.py
    • utils
      • __init__.py
      • helpers.py
      • metrics.py
      • visualization.py
      • logging.py
  • data
    • raw
    • processed
    • external
    • interim
  • models
    • checkpoints
    • pretrained
    • exported
  • notebooks
    • 01_data_exploration.ipynb
    • 02_model_development.ipynb
    • 03_evaluation.ipynb
  • tests
    • __init__.py
    • test_models.py
    • test_data.py
    • test_training.py
  • configs
    • base_config.yaml
    • resnet_config.yaml
    • transformer_config.yaml
  • scripts
    • download_data.py
    • preprocess_data.py
  • Dockerfile
  • requirements.txt
  • setup.py
  • README.md

Directory Structure Explanation

src/export.py

Script for exporting trained models to different formats (ONNX, TorchScript).

src/models/layers.py

Custom PyTorch layers and building blocks for neural networks.

src/data/augmentation.py

Advanced data augmentation techniques and custom transforms.

src/training/

Training infrastructure including trainers, optimizers, and schedulers.

src/training/trainer.py

Training class that encapsulates the training loop and validation logic.

src/utils/visualization.py

Functions for visualizing training progress, model predictions, and data.

models/exported/

Models exported to production formats like ONNX or TorchScript.

tests/

Unit tests for models, data processing, and training components.

configs/

YAML configuration files for different model architectures and experiments.

scripts/

Utility scripts for data preparation, preprocessing, and automation.