2025 USA-NA-AIO Round 2, Problem 3, Part 16

AI Olympiads 2025 USA-NA-AIO Round 2
1

Part 16 (5 points, coding task)

In this part, you are asked to train your model.

  1. Set the number of epochs as 100.

  2. Do training on GPU.

  3. For every epoch, print the average loss per sample in this epoch.

  4. You may use tqdm to track your progress and help you manage your time.

2 
### WRITE YOUR SOLUTION HERE ###

num_epochs = 100
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model_CLIP.to(device)

for epoch in tqdm(range(num_epochs)):
    model_CLIP.train()
    optimizer.zero_grad()
    loss_cum = 0
    for image_batch, token_id_batch, attention_mask_batch in CLIP_dataloader:
        image_batch = image_batch.to(device)
        token_id_batch = token_id_batch.to(device)
        attention_mask_batch = attention_mask_batch.to(device)

        image_embedding, text_embedding = model_CLIP(image_batch, token_id_batch, attention_mask_batch)
        loss = CLIP_loss_fn(image_embedding, text_embedding)
        loss.backward()
        optimizer.step()

        loss_cum += loss.item() * image_batch.shape[0]

    loss = loss_cum / len(CLIP_dateset)
    print(f'Epoch {epoch}, Loss: {loss}')

""" END OF THIS PART """