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