Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| """ Class average finetuning functions. Before using any of these finetuning | |
| functions, ensure that the model is set up with nb_classes=2. | |
| """ | |
| from __future__ import print_function | |
| import uuid | |
| from time import sleep | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from torchmoji.global_variables import ( | |
| FINETUNING_METHODS, | |
| WEIGHTS_DIR) | |
| from torchmoji.finetuning import ( | |
| freeze_layers, | |
| get_data_loader, | |
| fit_model, | |
| train_by_chain_thaw, | |
| find_f1_threshold) | |
| def relabel(y, current_label_nr, nb_classes): | |
| """ Makes a binary classification for a specific class in a | |
| multi-class dataset. | |
| # Arguments: | |
| y: Outputs to be relabelled. | |
| current_label_nr: Current label number. | |
| nb_classes: Total number of classes. | |
| # Returns: | |
| Relabelled outputs of a given multi-class dataset into a binary | |
| classification dataset. | |
| """ | |
| # Handling binary classification | |
| if nb_classes == 2 and len(y.shape) == 1: | |
| return y | |
| y_new = np.zeros(len(y)) | |
| y_cut = y[:, current_label_nr] | |
| label_pos = np.where(y_cut == 1)[0] | |
| y_new[label_pos] = 1 | |
| return y_new | |
| def class_avg_finetune(model, texts, labels, nb_classes, batch_size, | |
| method, epoch_size=5000, nb_epochs=1000, embed_l2=1E-6, | |
| verbose=True): | |
| """ Compiles and finetunes the given model. | |
| # Arguments: | |
| model: Model to be finetuned | |
| texts: List of three lists, containing tokenized inputs for training, | |
| validation and testing (in that order). | |
| labels: List of three lists, containing labels for training, | |
| validation and testing (in that order). | |
| nb_classes: Number of classes in the dataset. | |
| batch_size: Batch size. | |
| method: Finetuning method to be used. For available methods, see | |
| FINETUNING_METHODS in global_variables.py. Note that the model | |
| should be defined accordingly (see docstring for torchmoji_transfer()) | |
| epoch_size: Number of samples in an epoch. | |
| nb_epochs: Number of epochs. Doesn't matter much as early stopping is used. | |
| embed_l2: L2 regularization for the embedding layer. | |
| verbose: Verbosity flag. | |
| # Returns: | |
| Model after finetuning, | |
| score after finetuning using the class average F1 metric. | |
| """ | |
| if method not in FINETUNING_METHODS: | |
| raise ValueError('ERROR (class_avg_tune_trainable): ' | |
| 'Invalid method parameter. ' | |
| 'Available options: {}'.format(FINETUNING_METHODS)) | |
| (X_train, y_train) = (texts[0], labels[0]) | |
| (X_val, y_val) = (texts[1], labels[1]) | |
| (X_test, y_test) = (texts[2], labels[2]) | |
| checkpoint_path = '{}/torchmoji-checkpoint-{}.bin' \ | |
| .format(WEIGHTS_DIR, str(uuid.uuid4())) | |
| f1_init_path = '{}/torchmoji-f1-init-{}.bin' \ | |
| .format(WEIGHTS_DIR, str(uuid.uuid4())) | |
| if method in ['last', 'new']: | |
| lr = 0.001 | |
| elif method in ['full', 'chain-thaw']: | |
| lr = 0.0001 | |
| loss_op = nn.BCEWithLogitsLoss() | |
| # Freeze layers if using last | |
| if method == 'last': | |
| model = freeze_layers(model, unfrozen_keyword='output_layer') | |
| # Define optimizer, for chain-thaw we define it later (after freezing) | |
| if method == 'last': | |
| adam = optim.Adam((p for p in model.parameters() if p.requires_grad), lr=lr) | |
| elif method in ['full', 'new']: | |
| # Add L2 regulation on embeddings only | |
| special_params = [id(p) for p in model.embed.parameters()] | |
| base_params = [p for p in model.parameters() if id(p) not in special_params and p.requires_grad] | |
| embed_parameters = [p for p in model.parameters() if id(p) in special_params and p.requires_grad] | |
| adam = optim.Adam([ | |
| {'params': base_params}, | |
| {'params': embed_parameters, 'weight_decay': embed_l2}, | |
| ], lr=lr) | |
| # Training | |
| if verbose: | |
| print('Method: {}'.format(method)) | |
| print('Classes: {}'.format(nb_classes)) | |
| if method == 'chain-thaw': | |
| result = class_avg_chainthaw(model, nb_classes=nb_classes, | |
| loss_op=loss_op, | |
| train=(X_train, y_train), | |
| val=(X_val, y_val), | |
| test=(X_test, y_test), | |
| batch_size=batch_size, | |
| epoch_size=epoch_size, | |
| nb_epochs=nb_epochs, | |
| checkpoint_weight_path=checkpoint_path, | |
| f1_init_weight_path=f1_init_path, | |
| verbose=verbose) | |
| else: | |
| result = class_avg_tune_trainable(model, nb_classes=nb_classes, | |
| loss_op=loss_op, | |
| optim_op=adam, | |
| train=(X_train, y_train), | |
| val=(X_val, y_val), | |
| test=(X_test, y_test), | |
| epoch_size=epoch_size, | |
| nb_epochs=nb_epochs, | |
| batch_size=batch_size, | |
| init_weight_path=f1_init_path, | |
| checkpoint_weight_path=checkpoint_path, | |
| verbose=verbose) | |
| return model, result | |
| def prepare_labels(y_train, y_val, y_test, iter_i, nb_classes): | |
| # Relabel into binary classification | |
| y_train_new = relabel(y_train, iter_i, nb_classes) | |
| y_val_new = relabel(y_val, iter_i, nb_classes) | |
| y_test_new = relabel(y_test, iter_i, nb_classes) | |
| return y_train_new, y_val_new, y_test_new | |
| def prepare_generators(X_train, y_train_new, X_val, y_val_new, batch_size, epoch_size): | |
| # Create sample generators | |
| # Make a fixed validation set to avoid fluctuations in validation | |
| train_gen = get_data_loader(X_train, y_train_new, batch_size, | |
| extended_batch_sampler=True) | |
| val_gen = get_data_loader(X_val, y_val_new, epoch_size, | |
| extended_batch_sampler=True) | |
| X_val_resamp, y_val_resamp = next(iter(val_gen)) | |
| return train_gen, X_val_resamp, y_val_resamp | |
| def class_avg_tune_trainable(model, nb_classes, loss_op, optim_op, train, val, test, | |
| epoch_size, nb_epochs, batch_size, | |
| init_weight_path, checkpoint_weight_path, patience=5, | |
| verbose=True): | |
| """ Finetunes the given model using the F1 measure. | |
| # Arguments: | |
| model: Model to be finetuned. | |
| nb_classes: Number of classes in the given dataset. | |
| train: Training data, given as a tuple of (inputs, outputs) | |
| val: Validation data, given as a tuple of (inputs, outputs) | |
| test: Testing data, given as a tuple of (inputs, outputs) | |
| epoch_size: Number of samples in an epoch. | |
| nb_epochs: Number of epochs. | |
| batch_size: Batch size. | |
| init_weight_path: Filepath where weights will be initially saved before | |
| training each class. This file will be rewritten by the function. | |
| checkpoint_weight_path: Filepath where weights will be checkpointed to | |
| during training. This file will be rewritten by the function. | |
| verbose: Verbosity flag. | |
| # Returns: | |
| F1 score of the trained model | |
| """ | |
| total_f1 = 0 | |
| nb_iter = nb_classes if nb_classes > 2 else 1 | |
| # Unpack args | |
| X_train, y_train = train | |
| X_val, y_val = val | |
| X_test, y_test = test | |
| # Save and reload initial weights after running for | |
| # each class to avoid learning across classes | |
| torch.save(model.state_dict(), init_weight_path) | |
| for i in range(nb_iter): | |
| if verbose: | |
| print('Iteration number {}/{}'.format(i+1, nb_iter)) | |
| model.load_state_dict(torch.load(init_weight_path)) | |
| y_train_new, y_val_new, y_test_new = prepare_labels(y_train, y_val, | |
| y_test, i, nb_classes) | |
| train_gen, X_val_resamp, y_val_resamp = \ | |
| prepare_generators(X_train, y_train_new, X_val, y_val_new, | |
| batch_size, epoch_size) | |
| if verbose: | |
| print("Training..") | |
| fit_model(model, loss_op, optim_op, train_gen, [(X_val_resamp, y_val_resamp)], | |
| nb_epochs, checkpoint_weight_path, patience, verbose=0) | |
| # Reload the best weights found to avoid overfitting | |
| # Wait a bit to allow proper closing of weights file | |
| sleep(1) | |
| model.load_state_dict(torch.load(checkpoint_weight_path)) | |
| # Evaluate | |
| y_pred_val = model(X_val).cpu().numpy() | |
| y_pred_test = model(X_test).cpu().numpy() | |
| f1_test, best_t = find_f1_threshold(y_val_new, y_pred_val, | |
| y_test_new, y_pred_test) | |
| if verbose: | |
| print('f1_test: {}'.format(f1_test)) | |
| print('best_t: {}'.format(best_t)) | |
| total_f1 += f1_test | |
| return total_f1 / nb_iter | |
| def class_avg_chainthaw(model, nb_classes, loss_op, train, val, test, batch_size, | |
| epoch_size, nb_epochs, checkpoint_weight_path, | |
| f1_init_weight_path, patience=5, | |
| initial_lr=0.001, next_lr=0.0001, verbose=True): | |
| """ Finetunes given model using chain-thaw and evaluates using F1. | |
| For a dataset with multiple classes, the model is trained once for | |
| each class, relabeling those classes into a binary classification task. | |
| The result is an average of all F1 scores for each class. | |
| # Arguments: | |
| model: Model to be finetuned. | |
| nb_classes: Number of classes in the given dataset. | |
| train: Training data, given as a tuple of (inputs, outputs) | |
| val: Validation data, given as a tuple of (inputs, outputs) | |
| test: Testing data, given as a tuple of (inputs, outputs) | |
| batch_size: Batch size. | |
| loss: Loss function to be used during training. | |
| epoch_size: Number of samples in an epoch. | |
| nb_epochs: Number of epochs. | |
| checkpoint_weight_path: Filepath where weights will be checkpointed to | |
| during training. This file will be rewritten by the function. | |
| f1_init_weight_path: Filepath where weights will be saved to and | |
| reloaded from before training each class. This ensures that | |
| each class is trained independently. This file will be rewritten. | |
| initial_lr: Initial learning rate. Will only be used for the first | |
| training step (i.e. the softmax layer) | |
| next_lr: Learning rate for every subsequent step. | |
| seed: Random number generator seed. | |
| verbose: Verbosity flag. | |
| # Returns: | |
| Averaged F1 score. | |
| """ | |
| # Unpack args | |
| X_train, y_train = train | |
| X_val, y_val = val | |
| X_test, y_test = test | |
| total_f1 = 0 | |
| nb_iter = nb_classes if nb_classes > 2 else 1 | |
| torch.save(model.state_dict(), f1_init_weight_path) | |
| for i in range(nb_iter): | |
| if verbose: | |
| print('Iteration number {}/{}'.format(i+1, nb_iter)) | |
| model.load_state_dict(torch.load(f1_init_weight_path)) | |
| y_train_new, y_val_new, y_test_new = prepare_labels(y_train, y_val, | |
| y_test, i, nb_classes) | |
| train_gen, X_val_resamp, y_val_resamp = \ | |
| prepare_generators(X_train, y_train_new, X_val, y_val_new, | |
| batch_size, epoch_size) | |
| if verbose: | |
| print("Training..") | |
| # Train using chain-thaw | |
| train_by_chain_thaw(model=model, train_gen=train_gen, | |
| val_gen=[(X_val_resamp, y_val_resamp)], | |
| loss_op=loss_op, patience=patience, | |
| nb_epochs=nb_epochs, | |
| checkpoint_path=checkpoint_weight_path, | |
| initial_lr=initial_lr, next_lr=next_lr, | |
| verbose=verbose) | |
| # Evaluate | |
| y_pred_val = model(X_val).cpu().numpy() | |
| y_pred_test = model(X_test).cpu().numpy() | |
| f1_test, best_t = find_f1_threshold(y_val_new, y_pred_val, | |
| y_test_new, y_pred_test) | |
| if verbose: | |
| print('f1_test: {}'.format(f1_test)) | |
| print('best_t: {}'.format(best_t)) | |
| total_f1 += f1_test | |
| return total_f1 / nb_iter | |