CNNSupervisedTrainer_mnist_mnistClassifier_net.py 6.44 KB
Newer Older
Nicola Gatto's avatar
Nicola Gatto committed
1 2 3 4 5 6 7 8
import mxnet as mx
import logging
import numpy as np
import time
import os
import shutil
from mxnet import gluon, autograd, nd

9
class CNNSupervisedTrainer_mnist_mnistClassifier_net:
Sebastian N.'s avatar
Sebastian N. committed
10
    def __init__(self, data_loader, net_constructor):
Nicola Gatto's avatar
Nicola Gatto committed
11 12
        self._data_loader = data_loader
        self._net_creator = net_constructor
Sebastian N.'s avatar
Sebastian N. committed
13
        self._networks = {}
Nicola Gatto's avatar
Nicola Gatto committed
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

    def train(self, batch_size=64,
              num_epoch=10,
              eval_metric='acc',
              optimizer='adam',
              optimizer_params=(('learning_rate', 0.001),),
              load_checkpoint=True,
              context='gpu',
              checkpoint_period=5,
              normalize=True):
        if context == 'gpu':
            mx_context = mx.gpu()
        elif context == 'cpu':
            mx_context = mx.cpu()
        else:
            logging.error("Context argument is '" + context + "'. Only 'cpu' and 'gpu are valid arguments'.")

        if 'weight_decay' in optimizer_params:
            optimizer_params['wd'] = optimizer_params['weight_decay']
            del optimizer_params['weight_decay']
        if 'learning_rate_decay' in optimizer_params:
            min_learning_rate = 1e-08
            if 'learning_rate_minimum' in optimizer_params:
                min_learning_rate = optimizer_params['learning_rate_minimum']
                del optimizer_params['learning_rate_minimum']
            optimizer_params['lr_scheduler'] = mx.lr_scheduler.FactorScheduler(
                                                   optimizer_params['step_size'],
                                                   factor=optimizer_params['learning_rate_decay'],
                                                   stop_factor_lr=min_learning_rate)
            del optimizer_params['step_size']
            del optimizer_params['learning_rate_decay']


        train_iter, test_iter, data_mean, data_std = self._data_loader.load_data(batch_size)
Sebastian N.'s avatar
Sebastian N. committed
48 49 50 51 52

        if normalize:
            self._net_creator.construct(context=mx_context, data_mean=data_mean, data_std=data_std)
        else:
            self._net_creator.construct(context=mx_context)
Nicola Gatto's avatar
Nicola Gatto committed
53 54 55 56 57 58 59 60

        begin_epoch = 0
        if load_checkpoint:
            begin_epoch = self._net_creator.load(mx_context)
        else:
            if os.path.isdir(self._net_creator._model_dir_):
                shutil.rmtree(self._net_creator._model_dir_)

Sebastian N.'s avatar
Sebastian N. committed
61
        self._networks = self._net_creator.networks
Nicola Gatto's avatar
Nicola Gatto committed
62 63 64 65 66 67 68

        try:
            os.makedirs(self._net_creator._model_dir_)
        except OSError:
            if not os.path.isdir(self._net_creator._model_dir_):
                raise

Sebastian N.'s avatar
Sebastian N. committed
69
        trainers = [mx.gluon.Trainer(network.collect_params(), optimizer, optimizer_params) for network in self._networks.values()]
Nicola Gatto's avatar
Nicola Gatto committed
70

71 72
        loss_functions = {}

Sebastian N.'s avatar
Sebastian N. committed
73 74 75 76 77 78 79 80 81 82 83
        for network in self._networks.values():
            for output_name, last_layer in network.last_layers.items():
                if last_layer == 'softmax':
                    loss_functions[output_name] = mx.gluon.loss.SoftmaxCrossEntropyLoss()
                elif last_layer == 'sigmoid':
                    loss_functions[output_name] = mx.gluon.loss.SigmoidBinaryCrossEntropyLoss()
                elif last_layer == 'linear':
                    loss_functions[output_name] = mx.gluon.loss.L2Loss()
                else:
                    loss_functions[output_name] = mx.gluon.loss.L2Loss()
                    logging.warning("Invalid last layer, defaulting to L2 loss")
Nicola Gatto's avatar
Nicola Gatto committed
84 85 86 87 88 89 90

        speed_period = 50
        tic = None

        for epoch in range(begin_epoch, begin_epoch + num_epoch):
            train_iter.reset()
            for batch_i, batch in enumerate(train_iter):
91 92
                image_data = batch.data[0].as_in_context(mx_context)
                predictions_label = batch.label[0].as_in_context(mx_context)
93

Nicola Gatto's avatar
Nicola Gatto committed
94
                with autograd.record():
Sebastian N.'s avatar
Sebastian N. committed
95
                    predictions_output = self._networks[0](image_data)
96 97

                    loss = loss_functions['predictions'](predictions_output, predictions_label)
Nicola Gatto's avatar
Nicola Gatto committed
98 99

                loss.backward()
Sebastian N.'s avatar
Sebastian N. committed
100 101 102

                for trainer in trainers:
                    trainer.step(batch_size)
Nicola Gatto's avatar
Nicola Gatto committed
103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121

                if tic is None:
                    tic = time.time()
                else:
                    if batch_i % speed_period == 0:
                        try:
                            speed = speed_period * batch_size / (time.time() - tic)
                        except ZeroDivisionError:
                            speed = float("inf")

                        logging.info("Epoch[%d] Batch[%d] Speed: %.2f samples/sec" % (epoch, batch_i, speed))

                        tic = time.time()

            tic = None

            train_iter.reset()
            metric = mx.metric.create(eval_metric)
            for batch_i, batch in enumerate(train_iter):
122 123 124 125 126 127
                image_data = batch.data[0].as_in_context(mx_context)

                labels = [
                    batch.label[0].as_in_context(mx_context)
                ]

Sebastian N.'s avatar
Sebastian N. committed
128
                predictions_output = self._networks[0](image_data)
129

130 131 132 133 134
                predictions = [
                    mx.nd.argmax(predictions_output, axis=1)
                ]

                metric.update(preds=predictions, labels=labels)
Nicola Gatto's avatar
Nicola Gatto committed
135 136 137 138 139
            train_metric_score = metric.get()[1]

            test_iter.reset()
            metric = mx.metric.create(eval_metric)
            for batch_i, batch in enumerate(test_iter):
140 141 142 143 144 145
                image_data = batch.data[0].as_in_context(mx_context)

                labels = [
                    batch.label[0].as_in_context(mx_context)
                ]

Sebastian N.'s avatar
Sebastian N. committed
146
                predictions_output = self._networks[0](image_data)
147 148 149
                predictions = [
                    mx.nd.argmax(predictions_output, axis=1)
                ]
150

151
                metric.update(preds=predictions, labels=labels)
Nicola Gatto's avatar
Nicola Gatto committed
152 153 154 155 156
            test_metric_score = metric.get()[1]

            logging.info("Epoch[%d] Train: %f, Test: %f" % (epoch, train_metric_score, test_metric_score))

            if (epoch - begin_epoch) % checkpoint_period == 0:
Sebastian N.'s avatar
Sebastian N. committed
157 158
                for i, network in self._networks.items():
                    network.save_parameters(self.parameter_path(i) + '-' + str(epoch).zfill(4) + '.params')
Nicola Gatto's avatar
Nicola Gatto committed
159

Sebastian N.'s avatar
Sebastian N. committed
160 161 162
        for i, network in self._networks.items():
            network.save_parameters(self.parameter_path(i) + '-' + str(num_epoch + begin_epoch).zfill(4) + '.params')
            network.export(self.parameter_path(i) + '_newest', epoch=0)
Nicola Gatto's avatar
Nicola Gatto committed
163

Sebastian N.'s avatar
Sebastian N. committed
164 165
    def parameter_path(self, index):
        return self._net_creator._model_dir_ + self._net_creator._model_prefix_ + '_' + str(index)