train_imagenet.py

"""
Usage: 
    train_imagenet.py --model="mymodel.h5" --data="/data/ImageNetDir" --resume

Note:

   The idea here is to pre-train a network on imagenet, or some large corpus, and then transfer weights.
   I used https://github.com/mf1024/ImageNet-Datasets-Downloader.git to create a large dataset:
   `cd ImageNet-Datasets-Downloader
   python downloader.py -data_root /data/ImageNetData -number_of_classes 1000 -images_per_class 1000`
"""
import os
from docopt import docopt
from donkeycar.parts.keras import *
from donkeycar.utils import *
import numpy as np
import tensorflow as tf
from tensorflow.python import keras
from tensorflow.python.keras.layers import Dense
from tensorflow.python.keras.models import Model
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator


def train(args):

        batch_size = 256
        epochs = 50
        img_size=(120, 160)

        img_datagen = ImageDataGenerator(
                rescale=1./255,
                shear_range=0.2,
                zoom_range=0.2,
                horizontal_flip=True,
                validation_split=0.2)

        train_generator = img_datagen.flow_from_directory(
                args["--data"],
                target_size=img_size,
                batch_size=batch_size,
                class_mode='categorical',
                subset="training")

        validation_generator = img_datagen.flow_from_directory(
                args["--data"],
                target_size=img_size,
                batch_size=batch_size,
                class_mode='categorical',
                subset="validation")


        num_classes = train_generator.num_classes
        print("num_classes", num_classes)

        num_train = train_generator.n
        num_val = validation_generator.n
        steps_train = num_train // batch_size
        steps_val = num_val // batch_size
        print("steps_train", steps_train)
        print("steps_val", steps_val)

        # Start with the usual Donkeycar Keras Linear model.
        model = KerasLinear().model

        # Ignore the last 7 layers and tack on two new dense fully connceted layers
        # to handle the new class output for "imagenet"
        x = Dense(num_classes * 2, activation='relu', name="nfc_1")(model.layers[-7].output)
        x = Dense(num_classes, activation='softmax', name="nfc_2")(x)
        in_img = model.input 

        # Construct a new model with the layers
        model2 = Model(inputs=[in_img], outputs=[x])

        # Optionally resume training by loading the weights
        if args['--resume']:
                print("loading weights from old model")
                model2.load_weights(args['--model'], by_name=True)

        model2.compile(loss='categorical_crossentropy', metrics=['acc'], optimizer="adam")

        save_best = keras.callbacks.ModelCheckpoint(args['--model'])
        callbacks_list = [save_best]

        model2.fit_generator(
                train_generator,
                steps_per_epoch=steps_train,
                epochs=epochs,
                validation_data=validation_generator,
                validation_steps=steps_val,
                callbacks=callbacks_list)

def convert(args):
        # Load the imagenet trained weights into the Keras Linear Donkeycar model
        model = KerasLinear().model
        model.load_weights(args['--model'], by_name=True)
        model.save(args["--model"].replace(".h5", "_converted.h5"))


if __name__ == "__main__":
        args = docopt(__doc__)
        train(args)
        convert(args)