Build & train

PyTorch training pattern

Shows that the SageMaker lifecycle stays stable while model code changes.

Inputs

numeric feature tensors + binary labels

Outputs

trained state_dict + feature/preprocessing schema

Core explanation

SageMaker manages the training job, while PyTorch remains responsible for the neural-network training logic.

Mental model

What this means before worrying about syntax.

Your PyTorch script defines network, loss, optimizer and training loop; SageMaker manages the job infrastructure around it.

SageMaker manages infrastructure; PyTorch owns network/loss/optimizer/training loop.
Training mode and evaluation mode change layer behavior.
For rare-event models, operating metrics matter more than raw accuracy.

Lifecycle position

features -> DataLoader -> forward/loss/backprop -> state_dict -> model artifact

Implementation anatomy

The sequence you would actually build or configure.

1
Define nn.Module.
2
Create Dataset/DataLoader.
3
Forward pass.
4
Compute loss.
5
zero_grad -> backward -> optimizer.step.
6
Evaluate with model.eval()/no_grad.
7
Persist state_dict + architecture/config.

Objects / terms worth remembering

torch.nn.ModuleDataLoaderoptimizerModelTrainer

Concrete example

A real model-development or production flow.

20 engineered fraud features

↓

MLP 20 -> 64 -> 32 -> 1 logit

↓

BCEWithLogitsLoss

↓

sigmoid at inference gives fraud probability

↓

threshold tuned to alert-volume/risk tradeoff

Implementation lab

Concrete code, data manipulation and configuration for this lifecycle step.

Define the fraud MLP

python

import torch
from torch import nn

class FraudMLP(nn.Module):
    def __init__(self, n_features: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_features, 64),
            nn.ReLU(),
            nn.Dropout(0.20),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Dropout(0.10),
            nn.Linear(32, 1),
        )

    def forward(self, x):
        return self.net(x).squeeze(1)  # logits

Mini-batch training loop

python

model = FraudMLP(n_features=X_train.shape[1])
loss_fn = nn.BCEWithLogitsLoss(pos_weight=positive_class_weight)
optimizer = torch.optim.AdamW(
    model.parameters(), lr=1e-3, weight_decay=1e-4
)

for epoch in range(20):
    model.train()

    for x_batch, y_batch in train_loader:
        optimizer.zero_grad()
        logits = model(x_batch)
        loss = loss_fn(logits, y_batch.float())
        loss.backward()
        optimizer.step()

    model.eval()
    with torch.no_grad():
        val_logits = model(X_val_tensor)
        val_prob = torch.sigmoid(val_logits)

Save weights + preprocessing metadata

python

import json
import os
import torch

model_dir = os.environ["SM_MODEL_DIR"]

torch.save(model.state_dict(), f"{model_dir}/model.pt")

with open(f"{model_dir}/feature_schema.json", "w") as f:
    json.dump({
        "features": feature_columns,
        "means": scaler.mean_.tolist(),
        "scales": scaler.scale_.tolist(),
        "threshold": 0.63,
    }, f, indent=2)

Design decisions that matter

Weighted loss/sampling for imbalance.
PR-AUC/recall at constrained false-positive rate.
CPU vs GPU.
Regularization and early stopping.

Common failure modes

model.train() is a PyTorch mode switch, not the SageMaker training-job call.
Dropout/batchnorm differ between train/eval.
Accuracy can be misleading for rare fraud.

Production readiness check

Items to close before treating this step as production-ready.

✓ normalization artifact✓ architecture config✓ state_dict✓ imbalance treatment✓ threshold✓ eval mode✓ device compatibility