Training Guide

This guide covers launching classification and segmentation training jobs — both locally and on SageMaker.

Available Models

Classification

Model	Notes
ResNet50	General-purpose CNN, good baseline
VGG16	Classic architecture, larger memory footprint
EfficientNet	Best accuracy/compute trade-off
DenseNet121	Dense connections, good for small datasets
ViT	Vision Transformer, requires more data
Custom	Upload your own model_def.py

Segmentation

Model	Notes
SegResNet	MONAI’s residual encoder-decoder, fast training
SwinUNETR	Transformer-based, state-of-the-art for 3D
UNet	Classic encoder-decoder architecture

Data Format

Upload your training data to S3 in the following structure:

Classification

s3://your-bucket/dataset-name/
├── train/
│   ├── class_a/
│   │   ├── image001.dcm
│   │   └── image002.dcm
│   └── class_b/
│       ├── image003.dcm
│       └── image004.dcm
├── val/
│   ├── class_a/
│   └── class_b/
└── test/        (optional — used for evaluation)
    ├── class_a/
    └── class_b/

• Each subfolder name becomes a class label

• num_classes is derived automatically from the folder structure

• Supported file formats: DICOM (.dcm), PNG, JPG, NIfTI

Segmentation

s3://your-bucket/dataset-name/
├── train/
│   ├── subject_001/
│   │   ├── img.nii.gz
│   │   └── mask.nii.gz
│   └── subject_002/
│       ├── img.nii.gz
│       └── mask.nii.gz
└── val/
    ├── subject_003/
    └── subject_004/

• Each subject folder contains an image volume and a binary mask

• Output is always 1 channel (binary segmentation)

Launching a SageMaker Training Job

Step 1 — Open the Training tab

Navigate to http://localhost:8000 and click the Training tab.

Step 2 — Fill out the training form

Task Type and Dimension (2D/3D) appear at the top of the form — selecting these first determines which model architectures and configuration options are shown below.

Field	Description	Example
Task Type	classification or segmentation (top of form)	classification
Dimension	2D or 3D (top of form, determines model options)	2D
Data S3 URI	S3 path to your dataset root	s3://my-bucket/spine-xray/
Model	Architecture to train	DenseNet121
Epochs	Number of training epochs	50
Batch Size	Mini-batch size	16
Learning Rate	Initial learning rate	0.0001
Instance Type	SageMaker instance	ml.g5.xlarge
Training Image URI	ECR URI of the BYOC container	from .env.local
Role ARN	SageMaker execution role	from .env.local
MLflow Tracking URI	(Optional) MLflow server URL	http://your-mlflow:5000

Step 3 — Submit

Click Start Training. The job appears in the Active tab with status InProgress.

Step 4 — Monitor progress

• The job list auto-refreshes status from AWS for active jobs

• Once complete, the job moves to the Completed tab with billable seconds and estimated cost

MLflow Integration

To enable experiment tracking:

1. Set the MLFLOW_TRACKING_URI environment variable (in .env or the training form)

2. The training script automatically logs:

   • Hyperparameters (model, epochs, learning rate, batch size)

   • Per-epoch metrics (train loss, train accuracy, val accuracy, val loss)

   • Best model checkpoint as an artifact

   • result.json with training provenance

For SageMaker-managed MLflow tracking servers, use the dashboard’s presigned URL feature to get a temporary access URL.

Fine-Tuning from a Previous Job

To continue training from a completed job’s checkpoint:

1. Find the completed job’s model artifact S3 URI (e.g. s3://bucket/output/job-name/output/model.tar.gz)

2. In the training form, set Checkpoint S3 URI to that path

3. The training script will:

   • Download and extract the archive

   • Load best_model.pth (or model.pth) — restoring model weights and optimizer state

   • Set start_epoch to the checkpoint’s epoch + 1

   • Record the source URI in result.json under fine_tuned_from

If the checkpoint cannot be loaded, training falls back to starting from scratch with a warning.

Job Management

Tabs

Tab	Shows
Active	Jobs in InProgress or Stopping state
Completed	Successfully finished jobs
Failed	Jobs that errored out
Archived	Jobs you’ve hidden from the default view

Actions

Action	Description
Sync from AWS	Pull up to 100 jobs from AWS and merge into local store. Use ?name_contains= to filter.
Archive	Hide a job from the default list (still accessible in the Archived tab)
Comment	Attach a free-text note to any job
Delete Artifacts	Delete model.tar.gz from S3 and mark the job as artifacts_deleted
Remove	Delete a job from local tracking (only allowed when not active)
Fetch Result	Download result.json from the model artifact (reconstructs if missing)

Search

Use the search bar to filter jobs by name. The ?name_contains= query parameter also works with the Sync from AWS operation to pull only matching jobs.

Custom Models

To use a custom architecture:

Step 1 — Create your model code

Create a zip file containing at minimum a model_def.py with a get_model(num_classes) function:

# model_def.py
import torch.nn as nn
from torchvision import models

def get_model(num_classes):
    model = models.resnet18(pretrained=True)
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    return model

Step 2 — Upload the zip

curl -X POST "http://localhost:8000/api/training/models/upload-code?model_name=my_custom_model" \
  -F "file=@my_model.zip"

Or use the dashboard’s Upload Custom Model form.

Step 3 — Select in training form

Choose your custom model from the model dropdown when launching a job. The code is extracted to data/custom_models/<model_name>/ and bundled into the training container.

Local Training

For quick experiments without SageMaker:

1. Select Local as the training target

2. Fill in the same form fields (data path can be local or S3)

3. Training runs on the dashboard server with MLflow and TensorBoard tracking

4. Results appear in the local training history

Local training is useful for prototyping but limited by the server’s hardware. Use SageMaker for GPU-accelerated training on larger datasets.