Training a simple Vision Transformer¶

This guide walks you through a basic, single GPU model training workflow. You'll train a Vision Transformer (ViT) on a small image dataset using:

Hugging Face datasets to load the images,
timm to build the ViT model, and
PyTorch Lightning to run the training loop.

This guide assumes that you're comfortable with the basics of training image models, but new to FLAME. If you need an introduction to model training, we recommend starting with the PyTorch tutorial.

There are two ways to do this guide:

Interactively in a JupyterHub notebook
As a batch job

We recommend getting started with an interactive notebook for learning purposes, but switching to batch jobs as you scale.

Interactive notebook (start here)

Launch a GPU notebook, set up the environment, and train the ViT live with inline output.

Interactive guide
Batch Job

Save your script and submit it as a TrainJob that runs unattended. Builds on the notebook guide.

Batch guide

What you'll need first

A FLAME workspace you're a member of. If you don't have one yet, see Requesting access.
A GPU allocation for your group that includes the gh200 GPU type. Check with your PI if you're unsure.

The dataset you'll train on is public, so you need no Hugging Face account or token — it downloads anonymously.