Skip to content
GitLab
Switch branch/tag
Find file BlameHistoryPermalink
train.py 3.54 KiB
1 
"""
2 
Train our RNN on extracted features or images.
3 
"""
4 
from keras.callbacks import TensorBoard, ModelCheckpoint, EarlyStopping, CSVLogger
5 
from models import ResearchModels
6 
from data import DataSet
7 
import time
8 
import os.path
9
10 
def train(data_type, seq_length, model, saved_model=None,
11 
          class_limit=None, image_shape=None,
12 
          load_to_memory=False, batch_size=32, nb_epoch=100):
13 
    # Helper: Save the model.
14 
    checkpointer = ModelCheckpoint(
15 
        filepath=os.path.join('data', 'checkpoints', model + '-' + data_type + \
16 
            '.{epoch:03d}-{val_loss:.3f}.hdf5'),
17 
        verbose=1,
18 
        save_best_only=True)
19
20 
    # Helper: TensorBoard
21 
    tb = TensorBoard(log_dir=os.path.join('data', 'logs', model))
22
23 
    # Helper: Stop when we stop learning.
24 
    early_stopper = EarlyStopping(patience=5)
25
26 
    # Helper: Save results.
27 
    timestamp = time.time()
28 
    csv_logger = CSVLogger(os.path.join('data', 'logs', model + '-' + 'training-' + \
29 
        str(timestamp) + '.log'))
30
31 
    # Get the data and process it.
32 
    if image_shape is None:
33 
        data = DataSet(
34 
            seq_length=seq_length,
35 
            class_limit=class_limit
36 
        )
37 
    else:
38 
        data = DataSet(
39 
            seq_length=seq_length,
40 
            class_limit=class_limit,
41 
            image_shape=image_shape
42 
        )
43
44 
    # Get samples per epoch.
45 
    # Multiply by 0.7 to attempt to guess how much of data.data is the train set.
46 
    steps_per_epoch = (len(data.data) * 0.7) // batch_size
47
48 
    if load_to_memory:
49 
        # Get data.
50 
        X, y = data.get_all_sequences_in_memory('train', data_type)
51 
        X_test, y_test = data.get_all_sequences_in_memory('test', data_type)
52 
    else:
53 
        # Get generators.
54 
        generator = data.frame_generator(batch_size, 'train', data_type)
55 
        val_generator = data.frame_generator(batch_size, 'test', data_type)
56
57 
    # Get the model.
58 
    rm = ResearchModels(len(data.classes), model, seq_length, saved_model)
59
60 
    # Fit!
61 
    if load_to_memory:
62 
        # Use standard fit.
63 
        rm.model.fit(
64 
            X,
65 
            y,
66 
            batch_size=batch_size,
67 
            validation_data=(X_test, y_test),
68 
            verbose=1,
69 
            callbacks=[tb, early_stopper, csv_logger],
70 
            epochs=nb_epoch)
7172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
else: # Use fit generator. rm.model.fit_generator( generator=generator, steps_per_epoch=steps_per_epoch, epochs=nb_epoch, verbose=1, callbacks=[tb, early_stopper, csv_logger, checkpointer], validation_data=val_generator, validation_steps=40, workers=4) def main(): """These are the main training settings. Set each before running this file.""" # model can be one of lstm, lrcn, mlp, conv_3d, c3d model = 'lstm' saved_model = None # None or weights file class_limit = None # int, can be 1-101 or None seq_length = 40 load_to_memory = False # pre-load the sequences into memory batch_size = 32 nb_epoch = 1000 # Chose images or features and image shape based on network. if model in ['conv_3d', 'c3d', 'lrcn']: data_type = 'images' image_shape = (80, 80, 3) elif model in ['lstm', 'mlp']: data_type = 'features' image_shape = None else: raise ValueError("Invalid model. See train.py for options.") train(data_type, seq_length, model, saved_model=saved_model, class_limit=class_limit, image_shape=image_shape, load_to_memory=load_to_memory, batch_size=batch_size, nb_epoch=nb_epoch) if __name__ == '__main__': main()