Appendix H: Best Coding Practices and Style Guidelines for PyTorch

Writing clean, efficient, and maintainable PyTorch code is essential for building scalable machine learning projects. This appendix outlines recommended best practices, architectural guidelines, coding conventions, and workflow optimizations for PyTorch development.

1. Code Structure and Organization

A clean project structure improves readability and supports faster debugging.

1.1 Recommended Folder Structure

project/
│
├── data/                # datasets, raw and processed
├── models/              # model definitions (nn.Module files)
├── utils/               # helper functions (visualization, metrics)
├── configs/             # hyperparameters, settings
├── notebooks/           # experiments and prototyping
├── checkpoints/         # saved models
├── logs/                # training logs (TensorBoard, JSON)
└── train.py             # training script
└── inference.py         # testing/inference script

1.2 Keep Code Modular

Separate:

Data preprocessing
Model architecture
Training loop
Evaluation functions

Example:

from models.cnn import SimpleCNN
from utils.train import train_one_epoch
from utils.data import get_dataloader

2. Naming Conventions

2.1 Variables and Functions

snake_case for variables/functions
CamelCase for classes
Avoid single-letter names except math (e.g., x, y, z)

2.2 Good Examples

✔ train_one_epoch()
✔ compute_accuracy()
✔ CustomImageDataset
✔ learning_rate, batch_size

2.3 Bad Examples

✘ TrainFunction()
✘ acc()
✘ DS1, TmpVar, xysmp1

3. Best Practices for Creating Models

3.1 Use `nn.Module` Properly

class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(784, 10)

    def forward(self, x):
        return self.fc(x)

3.2 Avoid Logic Inside `forward()`

Only include:

Layer operations
Activations
Simple reshaping

❌ Avoid: file I/O, prints, condition-heavy debugging
✔ Use hooks for debugging

3.3 Initialize Weights Explicitly (Optional)

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)

model.apply(init_weights)

4. Device and Mixed Precision Management

4.1 Use a Consistent Device Assignment

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

4.2 Use Mixed Precision for Speed

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    output = model(x)
    loss = criterion(output, y)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

4.3 Use `.to(device)` Instead of `.cuda()`

Better compatibility across systems.

5. Best Practices for Training Loops

5.1 Keep Training Loop Clean & Minimal

def train_one_epoch(model, dataloader, optimizer, criterion, device):
    model.train()
    for x, y in dataloader:
        x, y = x.to(device), y.to(device)

        optimizer.zero_grad()
        out = model(x)
        loss = criterion(out, y)

        loss.backward()
        optimizer.step()

5.2 Use Learning Rate Schedulers

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
scheduler.step()

5.3 Log Training Progress Frequently

Avoid printing per batch (too slow).
Print:

Epoch-level metrics
Loss trends
LR changes

6. Data Handling Best Practices

6.1 Use `torchvision`, `torchtext`, `torchaudio`

Prefer existing datasets and transforms.

6.2 Create Efficient Custom Datasets

class MyDataset(Dataset):
    def __init__(self, files, transform=None):
        self.files = files
        self.transform = transform

    def __getitem__(self, idx):
        x = load_image(self.files[idx])
        if self.transform:
            x = self.transform(x)
        return x

6.3 Set `num_workers` for dataloaders

num_workers = 2 × number of CPU cores is ideal

6.4 Use Caching When Possible

✔ Cache tokenized text
✔ Preload audio spectrograms
✔ Save preprocessed tensors on disk

7. Style Guidelines for Clean PyTorch Code

7.1 Avoid Hard-Coding Values

Bad:

nn.Linear(784, 256)

Good:

input_dim = 784
hidden_dim = 256
nn.Linear(input_dim, hidden_dim)

**7.2 Avoid Using `*args` and `kwargs` in Models

It makes debugging harder.

7.3 Document Everything

Use docstrings:

def train(...):
    """Train the model for one epoch."""

7.4 Keep Functions Short

One function = one task.

8. Version Control and Experiment Tracking

8.1 Use Git for Version Control

Commit:

Model versions
Training loops
Experimental configurations

8.2 Track Experiments

Tools:

TensorBoard
MLflow
Weights & Biases
Neptune.ai

Track:

Loss curves
Hyperparameters
Activation distributions
Confusion matrices

9. Performance Optimization Tips

9.1 Prefer Vectorized Operations Over Loops

Avoid Python loops in tensor operations.

9.2 Move Static Tensors to Device Once

tensor = torch.tensor([...], device=device)

9.3 Use `with torch.no_grad()` During Evaluation

Saves memory and speeds up inference.

9.4 Use `pin_memory=True`

DataLoader(..., pin_memory=True)

9.5 Profile Slow Code

Use PyTorch Profiler:

import torch.profiler as profiler

10. Testing, Validation, and Reproducibility

10.1 Set Random Seeds

torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

10.2 Validate Every Component Independently

Test dataloader output
Test forward pass
Test loss computation
Test training on a small batch

10.3 Unit Test Custom Layers

Use pytest or simple assert statements.

11. Writing Reusable and Maintainable Code

11.1 Avoid Monolithic Scripts

Break into modules.

11.2 Use Config Files for Hyperparameters

learning_rate: 0.001
batch_size: 64
num_epochs: 30
model:
  type: resnet18

11.3 Save Model States Properly

torch.save(model.state_dict(), "model.pth")

Conclusion

This appendix provided a comprehensive guide to writing clean, efficient, scalable PyTorch code. Following these best practices ensures:

Faster debugging
Reproducible experiments
Better code readability
Easier collaboration
More stable training pipelines

With consistent coding style and modular design, PyTorch projects remain manageable even as complexity grows.