Träning med blandad precision i grundläggande PyTorch

Du ska använda flyttalstyper med låg precision för att snabba upp träningen av din språköversättningsmodell. Till exempel är 16-bitars flyttalstyper (float16) bara hälften så stora som sina 32-bitars motsvarigheter (float32). Det här påskyndar beräkningar av matrixmultiplikationer och konvolutioner. Kom ihåg att detta innebär att du skalar gradienter och castar operationer till 16-bitars flyttal.

Några objekt har förhandslästs in: dataset, model, dataloader och optimizer.

Den här övningen är en del av kursen

Effektiv AI-modellträning med PyTorch

Visa kurs

Övningsinstruktioner

Definiera en scaler för gradienterna med torch.amp.GradScaler innan loopen.
Använd torch.autocast som kontexthanterare i loopen för att casta operationer till 16-bitars flyttalstypen.
Skala förlusten i loopen och anropa .backward() för att skapa skalade gradienter.

Interaktiv övning med praktiskt arbete

Testa den här övningen genom att slutföra den här exempelkoden.

# Define a scaler for the gradients
scaler = torch.amp.____()
for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    # Casts operations to mixed precision
    with torch.____(device_type="cpu", dtype=torch.____):
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
    # Compute scaled gradients
    scaler.____(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

Redigera och kör kod