293 lines
11 KiB
Python
293 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import os
|
|
|
|
# import container_support as cs
|
|
import argparse
|
|
import json
|
|
|
|
import numpy as np
|
|
import re
|
|
import tensorflow as tf
|
|
import zipfile
|
|
# from tensorflow.keras import backend as K
|
|
from tensorflow.keras import callbacks
|
|
from tensorflow.keras.layers import Convolution2D
|
|
from tensorflow.keras.layers import Dropout, Flatten, Dense
|
|
from tensorflow.keras.layers import Input
|
|
from tensorflow.keras.models import Model
|
|
from tensorflow.keras.preprocessing.image import load_img, img_to_array
|
|
from tensorflow.python.client import device_lib
|
|
|
|
|
|
def linear_bin(a: float, N: int = 15, offset: int = 1, R: float = 2.0):
|
|
"""
|
|
create a bin of length N
|
|
map val A to range R
|
|
offset one hot bin by offset, commonly R/2
|
|
"""
|
|
a = a + offset
|
|
b = round(a / (R / (N - offset)))
|
|
arr = np.zeros(N)
|
|
b = clamp(b, 0, N - 1)
|
|
arr[int(b)] = 1
|
|
return arr
|
|
|
|
|
|
def clamp(n, min, max):
|
|
if n <= min:
|
|
return min
|
|
if n >= max:
|
|
return max
|
|
return n
|
|
|
|
|
|
def get_data(root_dir, filename):
|
|
print('load data from file ' + filename)
|
|
d = json.load(open(os.path.join(root_dir, filename)))
|
|
return [(d['user/angle']), root_dir, d['cam/image_array']]
|
|
|
|
|
|
numbers = re.compile(r'(\d+)')
|
|
|
|
|
|
def unzip_file(root, f):
|
|
zip_ref = zipfile.ZipFile(os.path.join(root, f), 'r')
|
|
zip_ref.extractall(root)
|
|
zip_ref.close()
|
|
|
|
|
|
def train(batch_size: int, slide_size: int, img_height: int, img_width: int, img_depth: int, horizon: int, drop: float):
|
|
# env = cs.TrainingEnvironment()
|
|
|
|
print(device_lib.list_local_devices())
|
|
os.system('mkdir -p logs')
|
|
|
|
# ### Loading the files ###
|
|
# ** You need to copy all your files to the directory where you are runing this notebook **
|
|
# ** into a folder named "data" **
|
|
|
|
data = []
|
|
|
|
for root, dirs, files in os.walk('/opt/ml/input/data/train'):
|
|
for f in files:
|
|
if f.endswith('.zip'):
|
|
unzip_file(root, f)
|
|
|
|
for root, dirs, files in os.walk('/opt/ml/input/data/train'):
|
|
data.extend(
|
|
[get_data(root, f) for f in sorted(files, key=str.lower) if f.startswith('record') and f.endswith('.json')])
|
|
|
|
# ### Loading throttle and angle ###
|
|
|
|
angle = [d[0] for d in data]
|
|
angle_array = np.array(angle)
|
|
|
|
# ### Loading images ###
|
|
if horizon > 0:
|
|
images = np.array([img_to_array(load_img(os.path.join(d[1], d[2])).crop((0, horizon, img_width, img_height))) for d in data], 'f')
|
|
else:
|
|
images = np.array( [img_to_array(load_img(os.path.join(d[1], d[2]))) for d in data], 'f')
|
|
|
|
# slide images vs orders
|
|
if slide_size > 0:
|
|
images = images[:len(images) - slide_size]
|
|
angle_array = angle_array[slide_size:]
|
|
|
|
# ### Start training ###
|
|
from datetime import datetime
|
|
logdir = '/opt/ml/model/logs/' + datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
logs = callbacks.TensorBoard(log_dir=logdir, histogram_freq=0, write_graph=True, write_images=True)
|
|
|
|
# Creates a file writer for the log directory.
|
|
# file_writer = tf.summary.create_file_writer(logdir)
|
|
|
|
# Using the file writer, log the reshaped image.
|
|
# with file_writer.as_default():
|
|
# # Don't forget to reshape.
|
|
# imgs = np.reshape(images[0:25], (-1, img_height, img_width, img_depth))
|
|
# tf.summary.image("25 training data examples", imgs, max_outputs=25, step=0)
|
|
|
|
save_best = callbacks.ModelCheckpoint('/opt/ml/model/model_cat', monitor='val_loss', verbose=1,
|
|
save_best_only=True, mode='min')
|
|
early_stop = callbacks.EarlyStopping(monitor='val_loss',
|
|
min_delta=.0005,
|
|
patience=5,
|
|
verbose=1,
|
|
mode='auto')
|
|
|
|
# categorical output of the angle
|
|
callbacks_list = [save_best, early_stop, logs]
|
|
|
|
angle_cat_array = np.array([linear_bin(float(a)) for a in angle_array])
|
|
|
|
model = default_model(input_shape=(img_height - horizon, img_width, img_depth), drop=drop)
|
|
#model = default_categorical(input_shape=(img_height - horizon, img_width, img_depth), drop=drop)
|
|
|
|
model.compile(optimizer='adam',
|
|
loss={'angle_out': 'categorical_crossentropy', },
|
|
loss_weights={'angle_out': 0.9})
|
|
model.fit({'img_in': images}, {'angle_out': angle_cat_array, }, batch_size=batch_size,
|
|
epochs=100, verbose=1, validation_split=0.2, shuffle=True, callbacks=callbacks_list)
|
|
|
|
# Save model for tensorflow using
|
|
model.save("/opt/ml/model/tfModel", save_format="tf")
|
|
|
|
def representative_dataset():
|
|
for d in tf.data.Dataset.from_tensor_slices(images).batch(1).take(100):
|
|
yield [tf.dtypes.cast(d, tf.float32)]
|
|
|
|
converter = tf.lite.TFLiteConverter.from_keras_model(model)
|
|
|
|
# full quantization for edgeTpu
|
|
# https://www.tensorflow.org/lite/performance/post_training_quantization#full_integer_quantization
|
|
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = representative_dataset
|
|
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
|
|
converter.inference_input_type = tf.uint8 # or tf.int8
|
|
converter.inference_output_type = tf.uint8 # or tf.int8
|
|
|
|
tflite_model = converter.convert()
|
|
|
|
# Save the model.
|
|
with open('/opt/ml/model/model_' + str(img_width) + 'x' + str(img_height) + 'h' + str(horizon) + '.tflite',
|
|
'wb') as f:
|
|
f.write(tflite_model)
|
|
|
|
|
|
def conv2d(filters, kernel, strides, layer_num, activation='relu'):
|
|
"""
|
|
Helper function to create a standard valid-padded convolutional layer
|
|
with square kernel and strides and unified naming convention
|
|
:param filters: channel dimension of the layer
|
|
:param kernel: creates (kernel, kernel) kernel matrix dimension
|
|
:param strides: creates (strides, strides) stride
|
|
:param layer_num: used in labelling the layer
|
|
:param activation: activation, defaults to relu
|
|
:return: tf.keras Convolution2D layer
|
|
"""
|
|
return Convolution2D(filters=filters,
|
|
kernel_size=(kernel, kernel),
|
|
strides=(strides, strides),
|
|
activation=activation,
|
|
name='conv2d_' + str(layer_num))
|
|
|
|
|
|
def core_cnn_layers(img_in: Input, img_height: int, img_width: int, drop: float, l4_stride: int = 1):
|
|
"""
|
|
Returns the core CNN layers that are shared among the different models,
|
|
like linear, imu, behavioural
|
|
:param img_in: input layer of network
|
|
:param drop: dropout rate
|
|
:param l4_stride: 4-th layer stride, default 1
|
|
:return: stack of CNN layers
|
|
"""
|
|
x = img_in
|
|
x = conv2d(img_height/5, 5, 2, 1)(x)
|
|
x = Dropout(drop)(x)
|
|
x = conv2d(img_width / 5, 5, 2, 2)(x)
|
|
x = Dropout(drop)(x)
|
|
x = conv2d(64, 5, 2, 3)(x)
|
|
x = Dropout(drop)(x)
|
|
x = conv2d(64, 3, l4_stride, 4)(x)
|
|
x = Dropout(drop)(x)
|
|
x = conv2d(64, 3, 1, 5)(x)
|
|
x = Dropout(drop)(x)
|
|
x = Flatten(name='flattened')(x)
|
|
return x
|
|
|
|
|
|
def default_model(input_shape, drop):
|
|
# First layer, input layer, Shape comes from camera.py resolution, RGB
|
|
img_in = Input(shape=input_shape, name='img_in')
|
|
kernel_size = 5
|
|
|
|
x = img_in
|
|
# 24 features, 5 pixel x 5 pixel kernel (convolution, feauture) window, 2wx2h stride, relu activation
|
|
x = Convolution2D(input_shape[1] / kernel_size, (kernel_size, kernel_size), strides=(2, 2), activation='relu')(x)
|
|
x = Dropout(drop)(x)
|
|
# 32 features, 5px5p kernel window, 2wx2h stride, relu activatiion
|
|
x = Convolution2D(input_shape[0] / kernel_size, (kernel_size, kernel_size), strides=(2, 2), activation='relu')(x)
|
|
x = Dropout(drop)(x)
|
|
# 64 features, 5px5p kernel window, 2wx2h stride, relu
|
|
x = Convolution2D(64, (kernel_size, kernel_size), strides=(2, 2), activation='relu')(x)
|
|
x = Dropout(drop)(x)
|
|
# 64 features, 3px3p kernel window, 2wx2h stride, relu
|
|
x = Convolution2D(64, (3, 3), strides=(2, 2), activation='relu')(x)
|
|
x = Dropout(drop)(x)
|
|
# 64 features, 3px3p kernel window, 1wx1h stride, relu
|
|
x = Convolution2D(64, (3, 3), strides=(1, 1), activation='relu')(x)
|
|
x = Dropout(drop)(x)
|
|
|
|
# Possibly add MaxPooling (will make it less sensitive to position in image).
|
|
# Camera angle fixed, so may not to be needed
|
|
|
|
x = Flatten(name='flattened')(x) # Flatten to 1D (Fully connected)
|
|
x = Dense(100, activation='relu')(x) # Classify the data into 100 features, make all negatives 0
|
|
x = Dropout(drop)(x)
|
|
x = Dense(50, activation='relu')(x)
|
|
x = Dropout(drop)(x)
|
|
# Connect every input with every output and output 15 hidden units. Use Softmax to give percentage.
|
|
# 15 categories and find best one based off percentage 0.0-1.0
|
|
angle_out = Dense(15, activation='softmax', name='angle_out')(x)
|
|
|
|
model = Model(inputs=[img_in], outputs=[angle_out])
|
|
|
|
return model
|
|
|
|
|
|
def default_n_linear(num_outputs, input_shape=(120, 160, 3), drop=0.2):
|
|
img_in = Input(shape=input_shape, name='img_in')
|
|
x = core_cnn_layers(img_in, img_width=input_shape[1], img_height=input_shape[0], drop=drop)
|
|
x = Dense(100, activation='relu', name='dense_1')(x)
|
|
x = Dropout(drop)(x)
|
|
x = Dense(50, activation='relu', name='dense_2')(x)
|
|
x = Dropout(drop)(x)
|
|
|
|
outputs = []
|
|
for i in range(num_outputs):
|
|
outputs.append(
|
|
Dense(1, activation='linear', name='n_outputs' + str(i))(x))
|
|
|
|
model = Model(inputs=[img_in], outputs=outputs, name='linear')
|
|
return model
|
|
|
|
|
|
def default_categorical(input_shape=(120, 160, 3), drop=0.2):
|
|
img_in = Input(shape=input_shape, name='img_in')
|
|
x = core_cnn_layers(img_in, img_width=input_shape[1], img_height=input_shape[0], drop=drop, l4_stride=2)
|
|
x = Dense(100, activation='relu', name="dense_1")(x)
|
|
x = Dropout(drop)(x)
|
|
x = Dense(50, activation='relu', name="dense_2")(x)
|
|
x = Dropout(drop)(x)
|
|
# Categorical output of the angle into 15 bins
|
|
angle_out = Dense(15, activation='softmax', name='angle_out')(x)
|
|
|
|
model = Model(inputs=[img_in], outputs=[angle_out],
|
|
name='categorical')
|
|
return model
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument("--slide_size", type=int, default=0)
|
|
parser.add_argument("--img_height", type=int, default=120)
|
|
parser.add_argument("--img_width", type=int, default=160)
|
|
parser.add_argument("--img_depth", type=int, default=3)
|
|
parser.add_argument("--horizon", type=int, default=0)
|
|
parser.add_argument("--batch_size", type=int, default=32)
|
|
parser.add_argument("--drop", type=float, default=0.2)
|
|
|
|
args = parser.parse_args()
|
|
params = vars(args)
|
|
train(
|
|
batch_size=params["batch_size"],
|
|
slide_size=params["slide_size"],
|
|
img_height=params["img_height"],
|
|
img_width=params["img_width"],
|
|
img_depth=params["img_depth"],
|
|
horizon=params["horizon"],
|
|
drop=params["drop"],
|
|
)
|