src/train.py

"""
Training the network
"""
import datetime
import logging
import random
import time
from typing import Sequence, Tuple

import torch

import dataloader
from model import Decoder, Encoder, EncoderDecoderModel

# logging INFO, WARNING, ERROR, CRITICAL, DEBUG
logging.basicConfig(level=logging.INFO)
logging.disable(level=10)


def train_network(
    model: torch.nn.Module,
    train_set: Sequence[Tuple[torch.tensor, torch.Tensor]],
    dev_set: Sequence[Tuple[torch.tensor, torch.Tensor]],
    epochs: int,
    clip: int = 1,
):
    """
    Train the EncoderDecoderModel network for a given number of epoch
    -----------
    Parameters
        model: torch.nn.Module
            EncoderDecoderModel defined in model.py
        train_set: Sequence[Tuple[torch.tensor, torch.tensor]]
            tuple of vectorized (text, summary) from the training set
        dev_set: Sequence[Tuple[torch.tensor, torch.tensor]]
            tuple of vectorized (text, summary) for the dev set
        epochs: int
            the number of epochs to train on
        clip: int
            no idea
    Return
        None
    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    print("Device check. You are using:", model.device)

    # with torch.no_grad():

    optim = torch.optim.Adam(model.parameters(), lr=0.01)

    print("Epoch\ttrain loss\tdev accuracy\tcompute time")

    for epoch_n in range(epochs):
        # Tell the model it's in train mode for layers designed to
        # behave differently in train or evaluation
        # https://stackoverflow.com/questions/51433378/what-does-model-train-do-in-pytorch
        model.train()

        # To get the computing time per epoch
        epoch_start_time = time.time()

        # To get the model accuracy per epoch
        epoch_loss = 0.0
        epoch_length = 0

        # Iterates over all the text, summary tuples
        for source, target in train_set:
            source = source.to(device)
            target = target.to(device)

            # DEBUG Block
            # logging.debug("TRAIN")
            # logging.debug(f"cuda available ? {torch.cuda.is_available()}")
            # logging.debug(f"Source sur cuda ? {source.is_cuda}")
            # logging.debug(f"Target sur cuda ? {target.is_cuda}")

            out = model(source).to(device)
            logging.debug(f"outputs = {out.shape}")

            target = torch.nn.functional.pad(
                target, (0, len(out) - len(target)), value=-100
            )

            # logging.debug(f"prediction : {vectoriser.decode(output_predictions)}")
            loss = torch.nn.functional.nll_loss(out, target).to(device)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
            optim.step()

            epoch_loss += loss.item()
            epoch_length += source.shape[0]

        # To check the model accuracy on new data
        dev_correct = 0
        dev_total = 0

        # Iterates over text, summary tuple from dev
        for source, target in dev_set:
            # We here want to evaluate the model
            # so we're switching to evaluation mode
            model.eval()

            source = source.to(device)
            target = target.to(device)

            # We compute the result
            output = model(source).to(device)

            output_dim = output.shape[-1]

            output = output[1:].view(-1, output_dim)
            logging.debug(f"dev output : {output.shape}")
            target = target[1:].view(-1)
            # To compare the output with the target,
            # they have to be of same length so we're
            # padding the target with -100 idx that will
            # be ignored by the nll_loss function
            target = torch.nn.functional.pad(
                target, (0, len(output) - len(target)), value=-100
            )
            dev_loss = torch.nn.functional.nll_loss(output, target)
            dev_correct += dev_loss.item()
            dev_total += source.shape[0]

        # Compute of the epoch training time
        epoch_compute_time = time.time() - epoch_start_time

        print(
            f"{epoch_n}\t{epoch_loss/epoch_length:.5}\t{abs(dev_correct/dev_total):.2%}\t\t{datetime.timedelta(seconds=epoch_compute_time)}"
        )


def predict(model, tokens: Sequence[str]) -> Sequence[str]:
    """Predict the POS for a tokenized sequence"""
    words_idx = vectoriser.encode(tokens).to(device)
    # Pas de calcul de gradient ici : c'est juste pour les prédictions
    with torch.no_grad():
        # equivalent to model(input) when called out of class
        out = model(words_idx).to(device)
    out_predictions = out.to(device)
    print(out_predictions)
    out_predictions = out_predictions.argmax(dim=-1)
    return vectoriser.decode(out_predictions)


if __name__ == "__main__":
    train_dataset = dataloader.Data("data/train_extract.jsonl")
    words = train_dataset.get_words()
    vectoriser = dataloader.Vectoriser(words)

    train_dataset = dataloader.Data("data/train_extract.jsonl", transform=vectoriser)
    dev_dataset = dataloader.Data("data/dev_extract.jsonl", transform=vectoriser)

    train_dataloader = torch.utils.data.DataLoader(
        train_dataset, batch_size=2, shuffle=True, collate_fn=dataloader.pad_collate
    )

    dev_dataloader = torch.utils.data.DataLoader(
        dev_dataset, batch_size=4, shuffle=True, collate_fn=dataloader.pad_collate
    )

    for i_batch, batch in enumerate(train_dataloader):
        print(i_batch, batch[0], batch[1])

    ### NEURAL NETWORK ###
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Device check. You are using:", device)

    ### RÉSEAU ENTRAÎNÉ ###
    # Pour s'assurer que les résultats seront les mêmes à chaque run du notebook
    torch.use_deterministic_algorithms(True)
    torch.manual_seed(0)
    random.seed(0)

    # On peut également entraîner encoder séparemment
    encoder = Encoder(len(vectoriser.idx_to_token) + 1, 256, 512, 0.5, device)
    decoder = Decoder(len(vectoriser.idx_to_token) + 1, 256, 512, 0.5, device)

    trained_classifier = EncoderDecoderModel(encoder, decoder, vectoriser, device).to(
        device
    )

    print(next(trained_classifier.parameters()).device)
    # print(train_dataset.is_cuda)

    train_network(
        trained_classifier,
        train_dataset,
        dev_dataset,
        2,
    )

    torch.save(trained_classifier.state_dict(), "model/model.pt")
    vectoriser.save("model/vocab.pkl")
    trained_classifier.config.to_json_file("config.json")

    print(f"test summary : {vectoriser.decode(dev_dataset[6][1])}")
    print(
        f"test prediction : {predict(trained_classifier, vectoriser.decode(dev_dataset[6][0]))}"
    )