File size: 25,946 Bytes
22d4f29
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
# -*- coding: utf-8 -*-
""" Finetuning functions for doing transfer learning to new datasets.
"""
from __future__ import print_function

import uuid
from time import sleep
from io import open

import math
import pickle
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import accuracy_score
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.sampler import BatchSampler, SequentialSampler
from torch.nn.utils import clip_grad_norm

from sklearn.metrics import f1_score

from torchmoji.global_variables import (FINETUNING_METHODS,
                                               FINETUNING_METRICS,
                                               WEIGHTS_DIR)
from torchmoji.tokenizer import tokenize
from torchmoji.sentence_tokenizer import SentenceTokenizer

try:
    unicode
    IS_PYTHON2 = True
except NameError:
    unicode = str
    IS_PYTHON2 = False


def load_benchmark(path, vocab, extend_with=0):
    """ Loads the given benchmark dataset.

        Tokenizes the texts using the provided vocabulary, extending it with
        words from the training dataset if extend_with > 0. Splits them into
        three lists: training, validation and testing (in that order).

        Also calculates the maximum length of the texts and the
        suggested batch_size.

    # Arguments:
        path: Path to the dataset to be loaded.
        vocab: Vocabulary to be used for tokenizing texts.
        extend_with: If > 0, the vocabulary will be extended with up to
            extend_with tokens from the training set before tokenizing.

    # Returns:
        A dictionary with the following fields:
            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).
            added: Number of tokens added to the vocabulary.
            batch_size: Batch size.
            maxlen: Maximum length of an input.
    """
    # Pre-processing dataset
    with open(path, 'rb') as dataset:
        if IS_PYTHON2:
            data = pickle.load(dataset)
        else:
            data = pickle.load(dataset, fix_imports=True)

    # Decode data
    try:
        texts = [unicode(x) for x in data['texts']]
    except UnicodeDecodeError:
        texts = [x.decode('utf-8') for x in data['texts']]

    # Extract labels
    labels = [x['label'] for x in data['info']]

    batch_size, maxlen = calculate_batchsize_maxlen(texts)

    st = SentenceTokenizer(vocab, maxlen)

    # Split up dataset. Extend the existing vocabulary with up to extend_with
    # tokens from the training dataset.
    texts, labels, added = st.split_train_val_test(texts,
                                                   labels,
                                                   [data['train_ind'],
                                                    data['val_ind'],
                                                    data['test_ind']],
                                                   extend_with=extend_with)
    return {'texts': texts,
            'labels': labels,
            'added': added,
            'batch_size': batch_size,
            'maxlen': maxlen}


def calculate_batchsize_maxlen(texts):
    """ Calculates the maximum length in the provided texts and a suitable
        batch size. Rounds up maxlen to the nearest multiple of ten.

    # Arguments:
        texts: List of inputs.

    # Returns:
        Batch size,
        max length
    """
    def roundup(x):
        return int(math.ceil(x / 10.0)) * 10

    # Calculate max length of sequences considered
    # Adjust batch_size accordingly to prevent GPU overflow
    lengths = [len(tokenize(t)) for t in texts]
    maxlen = roundup(np.percentile(lengths, 80.0))
    batch_size = 250 if maxlen <= 100 else 50
    return batch_size, maxlen



def freeze_layers(model, unfrozen_types=[], unfrozen_keyword=None):
    """ Freezes all layers in the given model, except for ones that are
        explicitly specified to not be frozen.

    # Arguments:
        model: Model whose layers should be modified.
        unfrozen_types: List of layer types which shouldn't be frozen.
        unfrozen_keyword: Name keywords of layers that shouldn't be frozen.

    # Returns:
        Model with the selected layers frozen.
    """
    # Get trainable modules
    trainable_modules = [(n, m) for n, m in model.named_children() if len([id(p) for p in m.parameters()]) !=  0]
    for name, module in trainable_modules:
        trainable = (any(typ in str(module) for typ in unfrozen_types) or
                     (unfrozen_keyword is not None and unfrozen_keyword.lower() in name.lower()))
        change_trainable(module, trainable, verbose=False)
    return model


def change_trainable(module, trainable, verbose=False):
    """ Helper method that freezes or unfreezes a given layer.

    # Arguments:
        module: Module to be modified.
        trainable: Whether the layer should be frozen or unfrozen.
        verbose: Verbosity flag.
    """
    
    if verbose: print('Changing MODULE', module, 'to trainable =', trainable)
    for name, param in module.named_parameters():
        if verbose: print('Setting weight', name, 'to trainable =', trainable)
        param.requires_grad = trainable

    if verbose:
        action = 'Unfroze' if trainable else 'Froze'
        if verbose: print("{} {}".format(action, module))


def find_f1_threshold(model, val_gen, test_gen, average='binary'):
    """ Choose a threshold for F1 based on the validation dataset
        (see https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4442797/
        for details on why to find another threshold than simply 0.5)

    # Arguments:
        model: pyTorch model
        val_gen: Validation set dataloader.
        test_gen: Testing set dataloader.

    # Returns:
        F1 score for the given data and
        the corresponding F1 threshold
    """
    thresholds = np.arange(0.01, 0.5, step=0.01)
    f1_scores = []

    model.eval()
    val_out = [(y, model(X)) for X, y in val_gen]
    y_val, y_pred_val = (list(t) for t in zip(*val_out))

    test_out = [(y, model(X)) for X, y in test_gen]
    y_test, y_pred_test = (list(t) for t in zip(*val_out))

    for t in thresholds:
        y_pred_val_ind = (y_pred_val > t)
        f1_val = f1_score(y_val, y_pred_val_ind, average=average)
        f1_scores.append(f1_val)

    best_t = thresholds[np.argmax(f1_scores)]
    y_pred_ind = (y_pred_test > best_t)
    f1_test = f1_score(y_test, y_pred_ind, average=average)
    return f1_test, best_t


def finetune(model, texts, labels, nb_classes, batch_size, method,
             metric='acc', epoch_size=5000, nb_epochs=1000, embed_l2=1E-6,
             verbose=1):
    """ Compiles and finetunes the given pytorch 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.
        metric: Evaluation metric to be used. For available metrics, see
            FINETUNING_METRICS in global_variables.py.
        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 provided metric.
    """

    if method not in FINETUNING_METHODS:
        raise ValueError('ERROR (finetune): Invalid method parameter. '
                         'Available options: {}'.format(FINETUNING_METHODS))
    if metric not in FINETUNING_METRICS:
        raise ValueError('ERROR (finetune): Invalid metric parameter. '
                         'Available options: {}'.format(FINETUNING_METRICS))

    train_gen = get_data_loader(texts[0], labels[0], batch_size,
                                extended_batch_sampler=True, epoch_size=epoch_size)
    val_gen = get_data_loader(texts[1], labels[1], batch_size,
                              extended_batch_sampler=False)
    test_gen = get_data_loader(texts[2], labels[2], batch_size,
                              extended_batch_sampler=False)

    checkpoint_path = '{}/torchmoji-checkpoint-{}.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() if nb_classes <= 2 \
         else nn.CrossEntropyLoss()

    # 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
        embed_params_id = [id(p) for p in model.embed.parameters()]
        output_layer_params_id = [id(p) for p in model.output_layer.parameters()]
        base_params = [p for p in model.parameters()
                       if id(p) not in embed_params_id and id(p) not in output_layer_params_id and p.requires_grad]
        embed_params = [p for p in model.parameters() if id(p) in embed_params_id and p.requires_grad]
        output_layer_params = [p for p in model.parameters() if id(p) in output_layer_params_id and p.requires_grad]
        adam = optim.Adam([
            {'params': base_params},
            {'params': embed_params, 'weight_decay': embed_l2},
            {'params': output_layer_params, 'lr': 0.001},
            ], lr=lr)

    # Training
    if verbose:
        print('Method:  {}'.format(method))
        print('Metric:  {}'.format(metric))
        print('Classes: {}'.format(nb_classes))

    if method == 'chain-thaw':
        result = chain_thaw(model, train_gen, val_gen, test_gen, nb_epochs, checkpoint_path, loss_op, embed_l2=embed_l2,
                            evaluate=metric, verbose=verbose)
    else:
        result = tune_trainable(model, loss_op, adam, train_gen, val_gen, test_gen, nb_epochs, checkpoint_path,
                                evaluate=metric, verbose=verbose)
    return model, result


def tune_trainable(model, loss_op, optim_op, train_gen, val_gen, test_gen,
                   nb_epochs, checkpoint_path, patience=5, evaluate='acc',
                   verbose=2):
    """ Finetunes the given model using the accuracy 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.
        checkpoint_weight_path: Filepath where weights will be checkpointed to
            during training. This file will be rewritten by the function.
        patience: Patience for callback methods.
        evaluate: Evaluation method to use. Can be 'acc' or 'weighted_f1'.
        verbose: Verbosity flag.

    # Returns:
        Accuracy of the trained model, ONLY if 'evaluate' is set.
    """
    if verbose:
        print("Trainable weights: {}".format([n for n, p in model.named_parameters() if p.requires_grad]))
        print("Training...")
        if evaluate == 'acc':
            print("Evaluation on test set prior training:", evaluate_using_acc(model, test_gen))
        elif evaluate == 'weighted_f1':
            print("Evaluation on test set prior training:", evaluate_using_weighted_f1(model, test_gen, val_gen))

    fit_model(model, loss_op, optim_op, train_gen, val_gen, nb_epochs, checkpoint_path, patience)

    # 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_path))
    if verbose >= 2:
        print("Loaded weights from {}".format(checkpoint_path))

    if evaluate == 'acc':
        return evaluate_using_acc(model, test_gen)
    elif evaluate == 'weighted_f1':
        return evaluate_using_weighted_f1(model, test_gen, val_gen)


def evaluate_using_weighted_f1(model, test_gen, val_gen):
    """ Evaluation function using macro weighted F1 score.

    # Arguments:
        model: Model to be evaluated.
        X_test: Inputs of the testing set.
        y_test: Outputs of the testing set.
        X_val: Inputs of the validation set.
        y_val: Outputs of the validation set.
        batch_size: Batch size.

    # Returns:
        Weighted F1 score of the given model.
    """
    # Evaluate on test and val data
    f1_test, _ = find_f1_threshold(model, test_gen, val_gen, average='weighted_f1')
    return f1_test


def evaluate_using_acc(model, test_gen):
    """ Evaluation function using accuracy.

    # Arguments:
        model: Model to be evaluated.
        test_gen: Testing data iterator (DataLoader)

    # Returns:
        Accuracy of the given model.
    """

    # Validate on test_data
    model.eval()
    accs = []
    for i, data in enumerate(test_gen):
        x, y = data
        outs = model(x)
        if model.nb_classes > 2:
            pred = torch.max(outs, 1)[1]
            acc = accuracy_score(y.squeeze().numpy(), pred.squeeze().numpy())
        else:
            pred = (outs >= 0).long()
            acc = (pred == y).double().sum() / len(pred)
        accs.append(acc)
    return np.mean(accs)


def chain_thaw(model, train_gen, val_gen, test_gen, nb_epochs, checkpoint_path, loss_op,
               patience=5, initial_lr=0.001, next_lr=0.0001, embed_l2=1E-6, evaluate='acc', verbose=1):
    """ Finetunes given model using chain-thaw and evaluates using accuracy.

    # Arguments:
        model: Model to be finetuned.
        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.
        initial_lr: Initial learning rate. Will only be used for the first
            training step (i.e. the output_layer layer)
        next_lr: Learning rate for every subsequent step.
        seed: Random number generator seed.
        verbose: Verbosity flag.
        evaluate: Evaluation method to use. Can be 'acc' or 'weighted_f1'.

    # Returns:
        Accuracy of the finetuned model.
    """
    if verbose:
        print('Training..')

    # Train using chain-thaw
    train_by_chain_thaw(model, train_gen, val_gen, loss_op, patience, nb_epochs, checkpoint_path,
                        initial_lr, next_lr, embed_l2, verbose)

    if evaluate == 'acc':
        return evaluate_using_acc(model, test_gen)
    elif evaluate == 'weighted_f1':
        return evaluate_using_weighted_f1(model, test_gen, val_gen)


def train_by_chain_thaw(model, train_gen, val_gen, loss_op, patience, nb_epochs, checkpoint_path,
                        initial_lr=0.001, next_lr=0.0001, embed_l2=1E-6, verbose=1):
    """ Finetunes model using the chain-thaw method.

    This is done as follows:
    1) Freeze every layer except the last (output_layer) layer and train it.
    2) Freeze every layer except the first layer and train it.
    3) Freeze every layer except the second etc., until the second last layer.
    4) Unfreeze all layers and train entire model.

    # Arguments:
        model: Model to be trained.
        train_gen: Training sample generator.
        val_data: Validation data.
        loss: Loss function to be used.
        finetuning_args: Training early stopping and checkpoint saving parameters
        epoch_size: Number of samples in an epoch.
        nb_epochs: Number of epochs.
        checkpoint_weight_path: Where weight checkpoints should be saved.
        batch_size: Batch size.
        initial_lr: Initial learning rate. Will only be used for the first
            training step (i.e. the output_layer layer)
        next_lr: Learning rate for every subsequent step.
        verbose: Verbosity flag.
    """
    # Get trainable layers
    layers = [m for m in model.children() if len([id(p) for p in m.parameters()]) !=  0]

    # Bring last layer to front
    layers.insert(0, layers.pop(len(layers) - 1))

    # Add None to the end to signify finetuning all layers
    layers.append(None)

    lr = None
    # Finetune each layer one by one and finetune all of them at once
    # at the end
    for layer in layers:
        if lr is None:
            lr = initial_lr
        elif lr == initial_lr:
            lr = next_lr

        # Freeze all except current layer
        for _layer in layers:
            if _layer is not None:
                trainable = _layer == layer or layer is None
                change_trainable(_layer, trainable=trainable, verbose=False)

        # Verify we froze the right layers
        for _layer in model.children():
            assert all(p.requires_grad == (_layer == layer) for p in _layer.parameters()) or layer is None

        if verbose:
            if layer is None:
                print('Finetuning all layers')
            else:
                print('Finetuning {}'.format(layer))

        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)

        fit_model(model, loss_op, adam, train_gen, val_gen, nb_epochs,
                  checkpoint_path, patience)

        # 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_path))
        if verbose >= 2:
            print("Loaded weights from {}".format(checkpoint_path))


def calc_loss(loss_op, pred, yv):
    if type(loss_op) is nn.CrossEntropyLoss:
        return loss_op(pred.squeeze(), yv.squeeze())
    else:
        return loss_op(pred.squeeze(), yv.squeeze().float())


def fit_model(model, loss_op, optim_op, train_gen, val_gen, epochs,
              checkpoint_path, patience):
    """ Analog to Keras fit_generator function.

    # Arguments:
        model: Model to be finetuned.
        loss_op: loss operation (BCEWithLogitsLoss or CrossEntropy for e.g.)
        optim_op: optimization operation (Adam e.g.)
        train_gen: Training data iterator (DataLoader)
        val_gen: Validation data iterator (DataLoader)
        epochs: Number of epochs.
        checkpoint_path: Filepath where weights will be checkpointed to
            during training. This file will be rewritten by the function.
        patience: Patience for callback methods.
        verbose: Verbosity flag.

    # Returns:
        Accuracy of the trained model, ONLY if 'evaluate' is set.
    """
    # Save original checkpoint
    torch.save(model.state_dict(), checkpoint_path)

    model.eval()
    best_loss = np.mean([calc_loss(loss_op, model(Variable(xv)), Variable(yv)).data.cpu().numpy()[0] for xv, yv in val_gen])
    print("original val loss", best_loss)

    epoch_without_impr = 0
    for epoch in range(epochs):
        for i, data in enumerate(train_gen):
            X_train, y_train = data
            X_train = Variable(X_train, requires_grad=False)
            y_train = Variable(y_train, requires_grad=False)
            model.train()
            optim_op.zero_grad()
            output = model(X_train)
            loss = calc_loss(loss_op, output, y_train)
            loss.backward()
            clip_grad_norm(model.parameters(), 1)
            optim_op.step()

            acc = evaluate_using_acc(model, [(X_train.data, y_train.data)])
            print("== Epoch", epoch, "step", i, "train loss", loss.data.cpu().numpy()[0], "train acc", acc)

        model.eval()
        acc = evaluate_using_acc(model, val_gen)
        print("val acc", acc)

        val_loss = np.mean([calc_loss(loss_op, model(Variable(xv)), Variable(yv)).data.cpu().numpy()[0] for xv, yv in val_gen])
        print("val loss", val_loss)
        if best_loss is not None and val_loss >= best_loss:
            epoch_without_impr += 1
            print('No improvement over previous best loss: ', best_loss)

        # Save checkpoint
        if best_loss is None or val_loss < best_loss:
            best_loss = val_loss
            torch.save(model.state_dict(), checkpoint_path)
            print('Saving model at', checkpoint_path)

        # Early stopping
        if epoch_without_impr >= patience:
            break

def get_data_loader(X_in, y_in, batch_size, extended_batch_sampler=True, epoch_size=25000, upsample=False, seed=42):
    """ Returns a dataloader that enables larger epochs on small datasets and
        has upsampling functionality.

    # Arguments:
        X_in: Inputs of the given dataset.
        y_in: Outputs of the given dataset.
        batch_size: Batch size.
        epoch_size: Number of samples in an epoch.
        upsample: Whether upsampling should be done. This flag should only be
            set on binary class problems.

    # Returns:
        DataLoader.
    """
    dataset = DeepMojiDataset(X_in, y_in)

    if extended_batch_sampler:
        batch_sampler = DeepMojiBatchSampler(y_in, batch_size, epoch_size=epoch_size, upsample=upsample, seed=seed)
    else:
        batch_sampler = BatchSampler(SequentialSampler(y_in), batch_size, drop_last=False)

    return DataLoader(dataset, batch_sampler=batch_sampler, num_workers=0)

class DeepMojiDataset(Dataset):
    """ A simple Dataset class.

    # Arguments:
        X_in: Inputs of the given dataset.
        y_in: Outputs of the given dataset.
    
    # __getitem__ output:
        (torch.LongTensor, torch.LongTensor)
    """
    def __init__(self, X_in, y_in):
        # Check if we have Torch.LongTensor inputs (assume Numpy array otherwise)
        if not isinstance(X_in, torch.LongTensor):
            X_in = torch.from_numpy(X_in.astype('int64')).long()
        if not isinstance(y_in, torch.LongTensor):
            y_in = torch.from_numpy(y_in.astype('int64')).long()

        self.X_in = torch.split(X_in, 1, dim=0)
        self.y_in = torch.split(y_in, 1, dim=0)

    def __len__(self):
        return len(self.X_in)

    def __getitem__(self, idx):
        return self.X_in[idx].squeeze(), self.y_in[idx].squeeze()

class DeepMojiBatchSampler(object):
    """A Batch sampler that enables larger epochs on small datasets and
        has upsampling functionality.

    # Arguments:
        y_in: Labels of the dataset.
        batch_size: Batch size.
        epoch_size: Number of samples in an epoch.
        upsample: Whether upsampling should be done. This flag should only be
            set on binary class problems.
        seed: Random number generator seed.

    # __iter__ output:
        iterator of lists (batches) of indices in the dataset
    """

    def __init__(self, y_in, batch_size, epoch_size, upsample, seed):
        self.batch_size = batch_size
        self.epoch_size = epoch_size
        self.upsample = upsample

        np.random.seed(seed)

        if upsample:
            # Should only be used on binary class problems
            assert len(y_in.shape) == 1
            neg = np.where(y_in.numpy() == 0)[0]
            pos = np.where(y_in.numpy() == 1)[0]
            assert epoch_size % 2 == 0
            samples_pr_class = int(epoch_size / 2)
        else:
            ind = range(len(y_in))

        if not upsample:
            # Randomly sample observations in a balanced way
            self.sample_ind = np.random.choice(ind, epoch_size, replace=True)
        else:
            # Randomly sample observations in a balanced way
            sample_neg = np.random.choice(neg, samples_pr_class, replace=True)
            sample_pos = np.random.choice(pos, samples_pr_class, replace=True)
            concat_ind = np.concatenate((sample_neg, sample_pos), axis=0)

            # Shuffle to avoid labels being in specific order
            # (all negative then positive)
            p = np.random.permutation(len(concat_ind))
            self.sample_ind = concat_ind[p]

            label_dist = np.mean(y_in.numpy()[self.sample_ind])
            assert(label_dist > 0.45)
            assert(label_dist < 0.55)

    def __iter__(self):
        # Hand-off data using batch_size
        for i in range(int(self.epoch_size/self.batch_size)):
            start = i * self.batch_size
            end = min(start + self.batch_size, self.epoch_size)
            yield self.sample_ind[start:end]

    def __len__(self):
        # Take care of the last (maybe incomplete) batch
        return (self.epoch_size + self.batch_size - 1) // self.batch_size