Interactive (Jupyter notebook)¶

In this guide you'll fine-tune a Vision Transformer to classify food photos, using Food-101 — a public dataset of 101,000 images across 101 dish categories. Rather than training from scratch, you'll start from a model already pretrained on ImageNet and adapt it to this new task — an approach called transfer learning that gets you a working classifier in minutes rather than hours.

Launching a GPU notebook¶

JupyterHub gives you a full JupyterLab notebook running in FLAME with a GPU attached, in your browser. No software to download.

1. Open https://jupyter.flamecluster.io and sign in (see JupyterHub for the login walkthrough).
2. You'll see a "Server Options" form. Select the following options:
   • Resources: choose Medium (8 CPU, 16 GB)
   • GPU Type: choose GH200
   • GPU Count: choose 1.
3. Click Start. The first launch can take a minute or two while the cluster boots your notebook pod and attaches your storage.

Confirm you have a GPU¶

Once your notebook server starts, you'll see the JupyterLab interface. Double-click on "Python 3 (ipykernel)" to open a new notebook.

Copy this code into the first cell and run it by pressing Shift + Enter:

import torch
print(torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("GPU:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "none")

You should see CUDA available: True and a GPU name containing "GH200". If torch.cuda.is_available() is False, you didn't request a GPU on the launch form. Go to the File → Hub Control Panel menu, stop your server using the "Stop my server" button on the Hub Control Panel, and start a new one with a GPU by following the steps above.

Install libraries¶

The default notebook image includes PyTorch, but not the other libraries we'll use. Install them by running this in a notebook cell:

%pip install lightning==2.6.5 timm==1.0.27 datasets==2.19.1 pillow==12.2.0

Train the model¶

Now let's build the training run. Paste each block below into its own notebook cell then run it with Shift + Enter.

We start by importing all the libraries we'll use:

import torch
import timm
import lightning as L
from torch.utils.data import DataLoader
from datasets import load_dataset

Now define the model. For this we use PyTorch Lightning, a framework on top of PyTorch that handles training boilerplace code (looping over batches, moving data and the model onto the GPU, scaling to multi-GPU, etc.). Instead of a hand-written training loop, you describe your model as a LightningModule subclass: you fill in what each part of training does, and Lightning manages the rest. We need to implement two methods:

• training_step (what happens in a single step)
• configure_optimizers (which optimizer to use)

The model itself is a ViT-Base from timm (vit_base_patch16_224), loaded with pretrained=True so we start from ImageNet-pretrained weights instead of random ones.

class LitViT(L.LightningModule):
    def __init__(self):
        super().__init__()
        # 101 classes matches the number in our training data
        self.model = timm.create_model("vit_base_patch16_224", pretrained=True, num_classes=101)
        self.criterion = torch.nn.CrossEntropyLoss()

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.model(x)
        loss = self.criterion(logits, y)
        accuracy = (logits.argmax(dim=1) == y).float().mean()

        # Display the loss and accuracy as the model trains
        self.log_dict(
            {"loss": loss, "acc": accuracy},
            prog_bar=True,
            on_step=True,
            on_epoch=False,
        )
        return loss

    def configure_optimizers(self):
        return torch.optim.AdamW(self.parameters(), lr=1e-5)

Now we create a model instance and use timm to create a transform to pre-process an image into the shape the ViT expects:

model = LitViT()

data_cfg = timm.data.resolve_model_data_config(model.model)
transform = timm.data.create_transform(**data_cfg, is_training=True)

This will take a few seconds and show a progress bar as timm downloads the pretrained weights for the ViT.

Now, download the Food 101 dataset:

train_ds = load_dataset("ethz/food101", split="train")

This will take a little while, as it's downloading several GB.

Now we wrap the dataset in a PyTorch DataLoader. The collate function converts the training data from the format it's stored in (a Pillow Image object) into the format the model expects (a tensor of shape (3, 224, 224)).

def collate(examples):
    images = torch.stack([transform(e["image"].convert("RGB")) for e in examples])
    labels = torch.tensor([e["label"] for e in examples])
    return images, labels

loader = DataLoader(
    train_ds,
    batch_size=1024,
    num_workers=8,
    collate_fn=collate,
    drop_last=True,
    shuffle=True,
)

Finally, set up a Lightning Trainer to run the training and kick it off by calling run, which moves everything onto the GPU and executes the training loop:

trainer = L.Trainer(
    precision="bf16-mixed",
    max_epochs=10,
    enable_checkpointing=True,
)
trainer.fit(model, loader)

A brief explanation of the arguments we pass to Trainer:

• precision: specifies the floating point precision. Lower precisions like 8- or 16-bit are often preferred for model training because they're faster and use less memory. We use bfloat16.
• max_epochs: stop training after a fixed number of passes over the dataset (in our case 10). In a more realstic training, this would be larger and you would use a stop condition instead of or in addition to max_epochs, but we're keeping it simple for this example.
• enable_checkpointing: Setting this to True saves the model's weights to disk as training progresses, so you can reload the trained model later rather than starting over if the training is interruped. This is an important part of designing for preemption.

Watch it train¶

Within a few seconds you'll see Lightning's progress bar, with the loss and accuracy live updating as the model trains (exact numbers will vary).

Epoch 0/4  ━━━╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 7/73 0:00:06 • 0:00:31 2.16it/s v_num: 28.000 loss: 4.590 acc: 0.031

The accuracy will start out low (less than 10%), but quickly climb, and should be around 80 or 90% by the time training finishes. Again, this is a small and simple example. Training this model on Food 101 is very quick and efficient because of the strong features in the pretrainined weights we used.

See it make a prediction¶

Now let's see the model we trained in action by running it on a handful of images from Food-101's validation split (which the model never saw during training, so this is a fair test of whether it generalizes).

First, we have to move the model back onto the GPU with .cuda() (Lightning automatically moves it back onto the CPU when training finishes), and put it in evaluation mode with .eval()¹.

model.cuda().eval();

Now extract the human readable class names from the dataset.

class_names = train_ds.features["label"].names

We can look at a few examples:

print(class_names[0:10])

Create a new transform to convert the data from Pillow images into tensors. This is identical to the transform we used for training except we set is_training to False, so timm doesn't do the random augmentations it does when is_training is True.

eval_transform = timm.data.create_transform(**data_cfg, is_training=False)

Now load the validation data.

val_ds = load_dataset("ethz/food101", split="validation")

And finally use the trained model to predict the class for a few examples.

for idx in [0, 5000, 12500, 18000, 24000]:
    example = val_ds[idx]
    img = eval_transform(example["image"].convert("RGB")).unsqueeze(0).cuda()
    with torch.no_grad():
        pred = model.model(img).argmax(dim=1).item()
    display(example["image"].resize((160, 160)))
    print(f"predicted: {class_names[pred]} | actual: {class_names[example['label']]}")

For each image you'll see the picture followed by the model's predicted label and the true label. After a few minutes of fine-tuning most should match. Explore a few more examples by changing the indices in the above code. Wrong predictions are often near-misses between visually similar dishes, exactly the kind of fine-grained distinction that more training steps help with.

Next: run the same training as a batch job¶

A notebook is ideal for prototyping, but it only runs while your session is open — close the tab and the run stops. Once a script works, you'll usually want to run it unattended: submit it, log out, and collect results later.

That's exactly what the Batch single-GPU training with TrainJob guide does. It builds directly on this one: you'll collect the training code above into a train.py file, then submit a TrainJob that runs it on a GPU in the background — queued fairly, surviving logout, and running to completion on its own. Continue with the batch guide.

Other directions¶

• Train for real, not just a benchmark. Raise max_steps, add a validation loop to track accuracy as you train, and save checkpoints to /workspace so your progress survives a restart. Borrowed GPUs can be reclaimed mid-run — see Designing for preemption.