La fonction PyTorch parfaite pour entraîner son modèle facilement !

Dans cet article on vous partage une fonction parfaite pour entraîner son modèle avec la librairie PyTorch !

PyTorch est une des librairies les plus utilisés pour faire du Deep Learning.

Cette librairie a comme particularité d’imposer au développeur de coder ses propres fonctions et classes pour entraîner son modèle.

Alors certes PyTorch a une approche plus complexe mais il permet plus de flexibilité, là où Keras nous simplifie la vie en la rendant plus standard.

Aujourd’hui je vous propose cette fonction prête à l’emploi pour entraîner facilement son modèle de classification avec PyTorch.

Entraîner son modèle PyTorch

La fonction d’entraînement

Voilà la fonction qui permet d’entraîner notre modèle tout en enregistrant la précision et l’erreur !

import time

def train(model, optimizer, loss_fn, train_dl, val_dl, epochs=100, device='cpu'):

    print('train() called: model=%s, opt=%s(lr=%f), epochs=%d, device=%s\n' % \
          (type(model).__name__, type(optimizer).__name__,
           optimizer.param_groups[0]['lr'], epochs, device))

    history = {} # Collects per-epoch loss and acc like Keras' fit().
    history['loss'] = []
    history['val_loss'] = []
    history['acc'] = []
    history['val_acc'] = []

    start_time_sec = time.time()

    for epoch in range(1, epochs+1):

        # --- TRAIN AND EVALUATE ON TRAINING SET -----------------------------
        model.train()
        train_loss         = 0.0
        num_train_correct  = 0
        num_train_examples = 0

        for batch in train_dl:

            optimizer.zero_grad()

            x    = batch[0].to(device)
            y    = batch[1].to(device)
            yhat = model(x)
            loss = loss_fn(yhat, y)

            loss.backward()
            optimizer.step()

            train_loss         += loss.data.item() * x.size(0)
            num_train_correct  += (torch.max(yhat, 1)[1] == y).sum().item()
            num_train_examples += x.shape[0]

        train_acc   = num_train_correct / num_train_examples
        train_loss  = train_loss / len(train_dl.dataset)


        # --- EVALUATE ON VALIDATION SET -------------------------------------
        model.eval()
        val_loss       = 0.0
        num_val_correct  = 0
        num_val_examples = 0

        for batch in val_dl:

            x    = batch[0].to(device)
            y    = batch[1].to(device)
            yhat = model(x)
            loss = loss_fn(yhat, y)

            val_loss         += loss.data.item() * x.size(0)
            num_val_correct  += (torch.max(yhat, 1)[1] == y).sum().item()
            num_val_examples += y.shape[0]

        val_acc  = num_val_correct / num_val_examples
        val_loss = val_loss / len(val_dl.dataset)


        if epoch == 1 or epoch % 10 == 0:
          print('Epoch %3d/%3d, train loss: %5.2f, train acc: %5.2f, val loss: %5.2f, val acc: %5.2f' % \
                (epoch, epochs, train_loss, train_acc, val_loss, val_acc))

        history['loss'].append(train_loss)
        history['val_loss'].append(val_loss)
        history['acc'].append(train_acc)
        history['val_acc'].append(val_acc)

    # END OF TRAINING LOOP


    end_time_sec       = time.time()
    total_time_sec     = end_time_sec - start_time_sec
    time_per_epoch_sec = total_time_sec / epochs
    print()
    print('Time total:     %5.2f sec' % (total_time_sec))
    print('Time per epoch: %5.2f sec' % (time_per_epoch_sec))

    return history

En fait cette fonction vient de StackOverflow, un très bon site pour les développeur qui se pose toutes sortes de questions.

Entraîner son modèle

Pour entraîner son modèle il faudra spécifier les variables :

• model – torch.nn.Module
• optimizer – torch.optim.Optimizer
• loss_fn – la loss function
• train_dl – DataLoader
• val_dl – DataLoader
• epochs – int
• device – ‘cpu’ ou ‘cuda’ pour le lancer avec un GPU

history = train(
    model = model,
    optimizer = optimizer,
    loss_fn = loss_fn,
    train_dl = train_loader,
    val_dl = val_loader,
    device='cuda')

On aura, en sortie, une variable history répertoriant les performances du modèle.

Afficher les performances

Pour afficher les performances du modèle, vous n’avez qu’à exécuter ce bout de code qui se sert de la librairie matplotlib pour tracer l’évolution du modèle lors de son entraînement :

import matplotlib.pyplot as plt

acc = history['acc']
val_acc = history['val_acc']
loss = history['loss']
val_loss = history['val_loss']
epochs = range(1, len(acc) + 1)

plt.plot(epochs, acc, 'b', label='Training acc')
plt.plot(epochs, val_acc, 'r', label='Validation acc')
plt.title('Training and validation accuracy')
plt.legend()
plt.figure()
plt.plot(epochs, loss, 'b', label='Training loss')
plt.plot(epochs, val_loss, 'r', label='Validation loss')
plt.title('Training and validation loss')
plt.legend()
plt.show()

Photo by Shifaaz shamoon on Unsplash