Is this the correct way to do ImageNet training on Torch XLA?

username: aftoul
cooked: <p>I’m new to TorchXLA and I decided to try training a model on ImageNet on Kaggle TPU. Following a few tutorials and adapting them to my needs, I came up with this code for the training loop:</p>
<pre><code class="lang-auto">def train(device_id, flags):
        device = xm.xla_device()
        rank = xm.get_local_ordinal()

        batch_size = BATCH_SIZE // xm.xrt_world_size()
        steps_per_epoch = DS_SIZE // (batch_size * xm.xrt_world_size())
        # Build datasets
        train_loader, val_loader = get_dataloaders(steps_per_epoch, batch_size, device)
        xm.rendezvous("loaded dataset")
    
        model = MyModel()
        model = model.to(device)
        xm.broadcast_master_param(model)
        model_params = sum(p.numel() for p in model.parameters())

        xm.master_print(f'Model parameters: {model_params:,d}')
    
        # Set up the optimizer
        optimizer = optim.AdamW(
            model.parameters(),
            lr=LR,
            weight_decay=WD
        )
    
        scheduler = cosine_scheduler_with_warmup(
            optimizer,
            total_epochs=EPOCHS,
            steps_per_epoch=math.ceil(steps_per_epoch / GRAD_ACC_STEPS),
            warmup_epochs=WARMUP_EPOCHS,
            initial_lr=0.01,
            end_lr=0.001
        )
    
        xm.rendezvous("loaded model and optimizer")

        start_epoch, train_history, test_history = SERIAL_EXEC.run(lambda: checkpoint_load(model, optimizer, scheduler))
        
        xm.rendezvous("loaded weights")
        xm.master_print("training begins")
    
        for epoch in range(start_epoch, start_epoch + EPOCHS):
            xm.master_print(f"Starting epoch {epoch + 1}")
            model.train()
            total_loss = torch.zeros((), device=device)
            local_total_batches = 0
            for step, (images, labels, _) in zip(range(1, steps_per_epoch+1), train_loader):
                images = images.to(device)
                labels = labels.to(device)
                
                # Get predictions
                outputs = model(images)

                # Compute distillation loss
                loss = loss_fn(outputs, labels)
                optimizer.zero_grad()
                loss.backward()
                # Do optimizer step with gradient clipping
                if step % GRAD_ACC_STEPS == 0 or step == steps_per_epoch:
                    xm.reduce_gradients(optimizer, pin_layout=True)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
                    optimizer.step()
                    scheduler.step()
                xm.mark_step()
                with torch.no_grad():
                    total_loss += loss * images.size(0)
                    local_total_batches += images.size(0)
    
            # Aggregate metrics across devices
            global_loss = xm.mesh_reduce("total_loss", total_loss.item(), sum)
            global_batches = xm.mesh_reduce("total_batches", local_total_batches, sum)
    
            average_loss = global_loss / global_batches
    
            xm.master_print(f"Epoch [{epoch+1}/{EPOCHS}], Training Loss: {average_loss:.4f}")
    
            # Evaluation loop
            val_loss, val_accuracy = eval_on_val(val_loader, model, device)
            xm.master_print(f"Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.4f}%")
    
            train_history.append([average_loss])
            test_history.append([val_loss, val_accuracy])
    
            xm.save({
                'model': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'scheduler': scheduler.state_dict(),
                'history': [train_history, test_history],
                'epoch': epoch + 1
            }, 'checkpoint.pth')
    
        xm.master_print("Training complete")
</code></pre>
<p>And the eval function:</p>
<pre><code class="lang-auto">def eval_on_val(val_loader, model, device):
    model.eval()
    val_loss = torch.zeros((), device=device)
    correct = torch.zeros((), device=device)
    total = 0

    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            val_loss += loss * labels.size(0)
            _, predicted = torch.max(outputs.data, 1)
            correct += (predicted == labels).float().sum()
            total += labels.size(0)

    global_val_loss = xm.mesh_reduce("val_loss", val_loss.item(), sum)
    global_correct = xm.mesh_reduce("val_correct", correct.item(), sum)
    global_total = xm.mesh_reduce("val_total", total, sum)

    avg_val_loss = global_val_loss / global_total
    val_accuracy = 100.0 * global_correct / global_total

    return avg_val_loss, val_accuracy
</code></pre>
<p>I have a few concerns with this code:</p>
<ul>
<li>Will the accuracy and loss be aggregated correctly?</li>
<li>Will there be an issue with the gradient accumulations?</li>
<li>Will the model need to be recompiled every epoch because of the eval function?</li>
<li>Did I use <code>xm.broadcast_master_param()</code> correctly?</li>
</ul>