Function train

language_translation/main.py:109–150

(model, train_dl, loss_fn, optim, special_symbols, opts)

# Train the model for 1 epoch
def train(model, train_dl, loss_fn, optim, special_symbols, opts):

    # Object for accumulating losses
    losses = 0

    # Put model into training mode
    model.train()
    for src, tgt in tqdm(train_dl, ascii=True):

        src = src.to(DEVICE)
        tgt = tgt.to(DEVICE)

        # We need to reshape the input slightly to fit into the transformer
        tgt_input = tgt[:-1, :]

        # Create masks
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input, special_symbols["<pad>"], DEVICE)

        # Pass into model, get probability over the vocab out
        logits = model(src, tgt_input, src_mask, tgt_mask,src_padding_mask, tgt_padding_mask, src_padding_mask)

        # Reset gradients before we try to compute the gradients over the loss
        optim.zero_grad()

        # Get original shape back
        tgt_out = tgt[1:, :]

        # Compute loss and gradient over that loss
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
        loss.backward()

        # Step weights
        optim.step()

        # Accumulate a running loss for reporting
        losses += loss.item()

        if opts.dry_run:
            break

    # Return the average loss
    return losses / len(list(train_dl))

# Check the model accuracy on the validation dataset
def validate(model, valid_dl, loss_fn, special_symbols):