Trainer.py 13.7 KB
Newer Older
1 2 3 4 5 6 7
#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
# @author: Tiago de Freitas Pereira <tiago.pereira@idiap.ch>
# @date: Tue 09 Aug 2016 15:25:22 CEST

import tensorflow as tf
from ..network import SequenceNetwork
8 9 10
import threading
import os
import bob.io.base
11
import bob.core
12
from ..analyzers import SoftmaxAnalizer
13
from tensorflow.core.framework import summary_pb2
14
import time
15
from bob.learn.tensorflow.datashuffler import OnlineSampling
16
from bob.learn.tensorflow.utils.session import Session
17
from .learning_rate import constant
18

19 20 21 22 23
#logger = bob.core.log.setup("bob.learn.tensorflow")

import logging
logger = logging.getLogger("bob.learn")

24

25 26 27 28 29 30
class Trainer(object):
    """
    One graph trainer.
    Use this trainer when your CNN is composed by one graph

    **Parameters**
31

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46
    architecture:
      The architecture that you want to run. Should be a :py:class`bob.learn.tensorflow.network.SequenceNetwork`

    optimizer:
      One of the tensorflow optimizers https://www.tensorflow.org/versions/r0.10/api_docs/python/train.html

    use_gpu: bool
      Use GPUs in the training

    loss: :py:class:`bob.learn.tensorflow.loss.BaseLoss`
      Loss function

    temp_dir: str
      The output directory

47
    learning_rate: `bob.learn.tensorflow.trainers.learning_rate`
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67
      Initial learning rate

    convergence_threshold:

    iterations: int
      Maximum number of iterations

    snapshot: int
      Will take a snapshot of the network at every `n` iterations

    prefetch: bool
      Use extra Threads to deal with the I/O

    model_from_file: str
      If you want to use a pretrained model

    analizer:
      Neural network analizer :py:mod:`bob.learn.tensorflow.analyzers`

    verbosity_level:
68 69

    """
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
70

71
    def __init__(self,
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
72
                 inputs,
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
73
                 graph,
74
                 optimizer=tf.train.AdamOptimizer(),
75 76
                 use_gpu=False,
                 loss=None,
77
                 temp_dir="cnn",
78

79
                 # Learning rate
80
                 learning_rate=None,
81

82
                 ###### training options ##########
83
                 convergence_threshold=0.01,
84
                 iterations=5000,
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
85 86
                 snapshot=500,
                 validation_snapshot=100,
87
                 prefetch=False,
88 89

                 ## Analizer
90
                 analizer=SoftmaxAnalizer(),
91

92 93 94
                 ### Pretrained model
                 model_from_file="",

95
                 verbosity_level=2):
96

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
97
        self.inputs = inputs
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
98
        self.graph = graph
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
99
        self.loss = loss
100 101 102 103

        if not isinstance(self.graph, tf.Tensor):
            raise ValueError("Expected a tf.Tensor as input for the keywork `graph`")

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
104 105
        self.predictor = self.loss(self.graph, inputs['label'])

106
        self.optimizer_class = optimizer
107
        self.use_gpu = use_gpu
108 109
        self.temp_dir = temp_dir

110 111 112 113
        if learning_rate is None and model_from_file == "":
            self.learning_rate = constant()
        else:
            self.learning_rate = learning_rate
114 115 116

        self.iterations = iterations
        self.snapshot = snapshot
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
117
        self.validation_snapshot = validation_snapshot
118
        self.convergence_threshold = convergence_threshold
119
        self.prefetch = prefetch
120

121 122 123 124 125 126
        # Training variables used in the fit
        self.optimizer = None
        self.training_graph = None
        self.train_data_shuffler = None
        self.summaries_train = None
        self.train_summary_writter = None
127
        self.thread_pool = None
128 129 130 131 132

        # Validation data
        self.validation_graph = None
        self.validation_summary_writter = None

133 134 135 136 137
        # Analizer
        self.analizer = analizer

        self.thread_pool = None
        self.enqueue_op = None
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
138
        self.global_step = None
139

140
        self.model_from_file = model_from_file
141
        self.session = None
142

143 144
        bob.core.log.set_verbosity_level(logger, verbosity_level)

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
145 146 147
    def __del__(self):
        tf.reset_default_graph()

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
148
    """
149
    def compute_graph(self, data_shuffler, prefetch=False, name="", training=True):
150 151 152 153 154
        Computes the graph for the trainer.

        ** Parameters **

            data_shuffler: Data shuffler
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
155
            prefetch: Uses prefetch
156
            name: Name of the graph
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
157
            training: Is it a training graph?
158 159

        # Defining place holders
160
        if prefetch:
161
            [placeholder_data, placeholder_labels] = data_shuffler.get_placeholders_forprefetch(name=name)
162 163 164 165 166 167 168

            # Defining a placeholder queue for prefetching
            queue = tf.FIFOQueue(capacity=10,
                                 dtypes=[tf.float32, tf.int64],
                                 shapes=[placeholder_data.get_shape().as_list()[1:], []])

            # Fetching the place holders from the queue
169
            self.enqueue_op = queue.enqueue_many([placeholder_data, placeholder_labels])
170 171 172 173 174 175 176
            feature_batch, label_batch = queue.dequeue_many(data_shuffler.batch_size)

            # Creating the architecture for train and validation
            if not isinstance(self.architecture, SequenceNetwork):
                raise ValueError("The variable `architecture` must be an instance of "
                                 "`bob.learn.tensorflow.network.SequenceNetwork`")
        else:
177
            [feature_batch, label_batch] = data_shuffler.get_placeholders(name=name)
178 179

        # Creating graphs and defining the loss
180
        network_graph = self.architecture.compute_graph(feature_batch, training=training)
181 182 183
        graph = self.loss(network_graph, label_batch)

        return graph
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
184
    """
185 186 187

    def get_feed_dict(self, data_shuffler):
        """
188
        Given a data shuffler prepared the dictionary to be injected in the graph
189 190 191 192

        ** Parameters **
            data_shuffler:

193
        """
194
        [data, labels] = data_shuffler.get_batch()
195

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
196 197
        feed_dict = {self.inputs['data']: data,
                     self.inputs['label']: labels}
198 199
        return feed_dict

200
    def fit(self, step):
201 202 203 204 205 206 207 208 209
        """
        Run one iteration (`forward` and `backward`)

        ** Parameters **
            session: Tensorflow session
            step: Iteration number

        """

210
        if self.prefetch:
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
211
            _, l, lr, summary = self.session.run([self.optimizer, self.predictor,
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
212
                                                  self.learning_rate, self.summaries_train])
213 214
        else:
            feed_dict = self.get_feed_dict(self.train_data_shuffler)
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
215
            _, l, lr, summary = self.session.run([self.optimizer, self.predictor,
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
216
                                                  self.learning_rate, self.summaries_train], feed_dict=feed_dict)
217

218 219
        logger.info("Loss training set step={0} = {1}".format(step, l))
        self.train_summary_writter.add_summary(summary, step)
220

221
    def create_general_summary(self):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
222
        """
223
        Creates a simple tensorboard summary with the value of the loss and learning rate
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
224
        """
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
225

226
        # Train summary
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
227
        tf.summary.scalar('loss', self.predictor)
228 229
        tf.summary.scalar('lr', self.learning_rate)
        return tf.summary.merge_all()
230

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
231
    """
232
    def start_thread(self):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
233

234 235 236 237
        Start pool of threads for pre-fetching

        **Parameters**
          session: Tensorflow session
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
238

239

240
        threads = []
241
        for n in range(3):
242
            t = threading.Thread(target=self.load_and_enqueue, args=())
243 244 245 246
            t.daemon = True  # thread will close when parent quits
            t.start()
            threads.append(t)
        return threads
247

248
    def load_and_enqueue(self):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
249

250
        Injecting data in the place holder queue
251 252 253

        **Parameters**
          session: Tensorflow session
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
254

255

256
        while not self.thread_pool.should_stop():
257 258
            [train_data, train_labels] = self.train_data_shuffler.get_batch()
            [train_placeholder_data, train_placeholder_labels] = self.train_data_shuffler.get_placeholders()
259

260 261 262
            feed_dict = {train_placeholder_data: train_data,
                         train_placeholder_labels: train_labels}

263
            self.session.run(self.enqueue_op, feed_dict=feed_dict)
264

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
265
    """
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
266

267
    def bootstrap_graphs_fromfile(self, train_data_shuffler, validation_data_shuffler):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
268 269
        """
        Bootstrap all the necessary data from file
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
270 271 272 273 274 275 276

         ** Parameters **
           session: Tensorflow session
           train_data_shuffler: Data shuffler for training
           validation_data_shuffler: Data shuffler for validation


Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
277
        """
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
278
        saver = self.architecture.load(self.model_from_file, clear_devices=False)
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
279 280 281 282 283 284 285 286

        # Loading training graph
        self.training_graph = tf.get_collection("training_graph")[0]

        # Loding other elements
        self.optimizer = tf.get_collection("optimizer")[0]
        self.learning_rate = tf.get_collection("learning_rate")[0]
        self.summaries_train = tf.get_collection("summaries_train")[0]
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
287
        self.global_step = tf.get_collection("global_step")[0]
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
288 289 290 291 292 293 294 295

        if validation_data_shuffler is not None:
            self.validation_graph = tf.get_collection("validation_graph")[0]

        self.bootstrap_placeholders_fromfile(train_data_shuffler, validation_data_shuffler)

        return saver

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313
    def compute_validation(self, data_shuffler, step):
        """
        Computes the loss in the validation set

        ** Parameters **
            session: Tensorflow session
            data_shuffler: The data shuffler to be used
            step: Iteration number

        """
        # Opening a new session for validation
        feed_dict = self.get_feed_dict(data_shuffler)
        l = self.session.run(self.predictor, feed_dict=feed_dict)

        summaries = [summary_pb2.Summary.Value(tag="loss", simple_value=float(l))]
        self.validation_summary_writter.add_summary(summary_pb2.Summary(value=summaries), step)
        logger.info("Loss VALIDATION set step={0} = {1}".format(step, l))

314 315
    def train(self, train_data_shuffler, validation_data_shuffler=None):
        """
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
316 317 318 319 320 321
        Train the network:

         ** Parameters **

           train_data_shuffler: Data shuffler for training
           validation_data_shuffler: Data shuffler for validation
322 323 324 325 326
        """

        # Creating directory
        bob.io.base.create_directories_safe(self.temp_dir)
        self.train_data_shuffler = train_data_shuffler
327

328
        logger.info("Initializing !!")
329
        self.session = Session.instance(new=True).session
330 331 332 333

        # Loading a pretrained model
        if self.model_from_file != "":
            logger.info("Loading pretrained model from {0}".format(self.model_from_file))
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
334
            saver = self.bootstrap_graphs_fromfile(train_data_shuffler, validation_data_shuffler)
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
335

336
            start_step = self.global_step.eval(session=self.session)
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
337

338
        else:
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
339
            start_step = 0
340 341

            # TODO: find an elegant way to provide this as a parameter of the trainer
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
342
            self.global_step = tf.Variable(0, trainable=False, name="global_step")
343
            tf.add_to_collection("global_step", self.global_step)
344 345 346

            # Preparing the optimizer
            self.optimizer_class._learning_rate = self.learning_rate
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
347
            self.optimizer = self.optimizer_class.minimize(self.predictor, global_step=self.global_step)
348 349 350
            tf.add_to_collection("optimizer", self.optimizer)
            tf.add_to_collection("learning_rate", self.learning_rate)

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
351 352 353
            self.summaries_train = self.create_general_summary()
            tf.add_to_collection("summaries_train", self.summaries_train)

354
            # Train summary
355
            tf.global_variables_initializer().run(session=self.session)
356 357

            # Original tensorflow saver object
358
            saver = tf.train.Saver(var_list=tf.global_variables())
359

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
360 361
        #if isinstance(train_data_shuffler, OnlineSampling):
        #    train_data_shuffler.set_feature_extractor(self.architecture, session=self.session)
362 363

        # Start a thread to enqueue data asynchronously, and hide I/O latency.
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
364 365 366 367
        #if self.prefetch:
        #    self.thread_pool = tf.train.Coordinator()
        #    tf.train.start_queue_runners(coord=self.thread_pool, sess=self.session)
        #    threads = self.start_thread()
368 369

        # TENSOR BOARD SUMMARY
370
        self.train_summary_writter = tf.summary.FileWriter(os.path.join(self.temp_dir, 'train'), self.session.graph)
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
371 372 373 374
        if validation_data_shuffler is not None:
            self.validation_summary_writter = tf.summary.FileWriter(os.path.join(self.temp_dir, 'validation'),
                                                                    self.session.graph)

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
375
        for step in range(start_step, self.iterations):
376
            start = time.time()
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
377
            self.fit(step)
378 379 380 381 382
            end = time.time()
            summary = summary_pb2.Summary.Value(tag="elapsed_time", simple_value=float(end-start))
            self.train_summary_writter.add_summary(summary_pb2.Summary(value=[summary]), step)

            # Running validation
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
383 384
            if validation_data_shuffler is not None and step % self.validation_snapshot == 0:
                self.compute_validation(validation_data_shuffler, step)
385

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
386 387 388
                #if self.analizer is not None:
                #    self.validation_summary_writter.add_summary(self.analizer(
                #         validation_data_shuffler, self.architecture, self.session), step)
389 390 391 392 393

            # Taking snapshot
            if step % self.snapshot == 0:
                logger.info("Taking snapshot")
                path = os.path.join(self.temp_dir, 'model_snapshot{0}.ckp'.format(step))
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
394 395
                saver.save(self.session, path)
                #self.architecture.save(saver, path)
396 397 398 399 400 401 402 403 404

        logger.info("Training finally finished")

        self.train_summary_writter.close()
        if validation_data_shuffler is not None:
            self.validation_summary_writter.close()

        # Saving the final network
        path = os.path.join(self.temp_dir, 'model.ckp')
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
405
        saver.save(self.session, path)
406 407 408 409

        if self.prefetch:
            # now they should definetely stop
            self.thread_pool.request_stop()
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
410
            #self.thread_pool.join(threads)