Trainer.py 13.9 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 98
        #if not isinstance(graph, SequenceNetwork):
        #    raise ValueError("`architecture` should be instance of `SequenceNetwork`")
99

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

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

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

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

120 121 122 123 124 125
        # 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
126
        self.thread_pool = None
127 128 129 130 131

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

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

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

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

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

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

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

        ** Parameters **

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

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

            # 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
168
            self.enqueue_op = queue.enqueue_many([placeholder_data, placeholder_labels])
169 170 171 172 173 174 175
            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:
176
            [feature_batch, label_batch] = data_shuffler.get_placeholders(name=name)
177 178

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

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

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

        ** Parameters **
            data_shuffler:

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

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

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

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

        """

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

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

220
    def create_general_summary(self):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
221
        """
222
        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
223
        """
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
224

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

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

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

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

238

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

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

249
        Injecting data in the place holder queue
250 251 252

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

254

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

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

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

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

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

         ** 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
276
        """
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
277
        saver = self.architecture.load(self.model_from_file, clear_devices=False)
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
278 279 280 281 282 283 284 285

        # 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
286
        self.global_step = tf.get_collection("global_step")[0]
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
287 288 289 290 291 292 293 294

        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
295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312
    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))

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

         ** Parameters **

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

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

327
        logger.info("Initializing !!")
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
328

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
329
        if not isinstance(self.graph, tf.Tensor):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
330
            raise NotImplemented("Not tensor still not implemented")
331

332
        self.session = Session.instance(new=True).session
333 334 335 336

        # 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
337
            saver = self.bootstrap_graphs_fromfile(train_data_shuffler, validation_data_shuffler)
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
338

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

341
        else:
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
342
            start_step = 0
343 344

            # 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
345
            self.global_step = tf.Variable(0, trainable=False, name="global_step")
346
            tf.add_to_collection("global_step", self.global_step)
347 348 349

            # Preparing the optimizer
            self.optimizer_class._learning_rate = self.learning_rate
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
350
            self.optimizer = self.optimizer_class.minimize(self.predictor, global_step=self.global_step)
351 352 353
            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
354 355 356
            self.summaries_train = self.create_general_summary()
            tf.add_to_collection("summaries_train", self.summaries_train)

357
            # Train summary
358
            tf.global_variables_initializer().run(session=self.session)
359 360

            # Original tensorflow saver object
361
            saver = tf.train.Saver(var_list=tf.global_variables())
362

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

        # Start a thread to enqueue data asynchronously, and hide I/O latency.
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
367 368 369 370
        #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()
371 372

        # TENSOR BOARD SUMMARY
373
        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
374 375 376 377
        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
378
        for step in range(start_step, self.iterations):
379
            start = time.time()
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
380
            self.fit(step)
381 382 383 384 385
            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
386 387
            if validation_data_shuffler is not None and step % self.validation_snapshot == 0:
                self.compute_validation(validation_data_shuffler, step)
388

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

            # 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
397 398
                saver.save(self.session, path)
                #self.architecture.save(saver, path)
399 400 401 402 403 404 405 406 407

        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
408
        saver.save(self.session, path)
409 410 411 412

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