robocars/robocar-training
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

feat: add option to build linear model

Browse Source
This commit is contained in:
Cyrille Nofficial 2022-06-09 09:36:46 +02:00
parent ad3eb16333
commit 91758cfffa
1 changed files with 22 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
src/tf_container/train.py
View File

@ -19,6 +19,9 @@ from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.python.client import device_lib
MODEL_CATEGORICAL = "categorical"
MODEL_LINEAR = "linear"
def linear_bin(a: float, N: int = 15, offset: int = 1, R: float = 2.0):
"""
@ -57,7 +60,7 @@ def unzip_file(root, f):
zip_ref.close()
def train(batch_size: int, slide_size: int, img_height: int, img_width: int, img_depth: int, horizon: int, drop: float):
def train(model_type: str, 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())
@ -108,7 +111,13 @@ def train(batch_size: int, slide_size: int, img_height: int, img_width: int, img
# 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,
model_filepath = '/opt/ml/model/model_other'
if model_type == MODEL_CATEGORICAL:
model_filepath = '/opt/ml/model/model_cat'
elif model_type == MODEL_LINEAR:
model_filepath = '/opt/ml/model/model_lin'
save_best = callbacks.ModelCheckpoint(model_filepath, monitor='val_loss', verbose=1,
save_best_only=True, mode='min')
early_stop = callbacks.EarlyStopping(monitor='val_loss',
min_delta=.0005,
@ -121,8 +130,8 @@ def train(batch_size: int, slide_size: int, img_height: int, img_width: int, img
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 = 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', },
@ -150,7 +159,7 @@ def train(batch_size: int, slide_size: int, img_height: int, img_width: int, img
tflite_model = converter.convert()
# Save the model.
with open('/opt/ml/model/model_' + str(img_width) + 'x' + str(img_height) + 'h' + str(horizon) + '.tflite',
with open('/opt/ml/model/model_' + model_type + '_' + str(img_width) + 'x' + str(img_height) + 'h' + str(horizon) + '.tflite',
'wb') as f:
f.write(tflite_model)
@ -177,6 +186,8 @@ def core_cnn_layers(img_in: Input, img_height: int, img_width: int, drop: float,
"""
Returns the core CNN layers that are shared among the different models,
like linear, imu, behavioural
:param img_width: image width
:param img_height: image height
:param img_in: input layer of network
:param drop: dropout rate
:param l4_stride: 4-th layer stride, default 1
@ -236,20 +247,16 @@ def default_model(input_shape, drop):
return model
def default_n_linear(num_outputs, input_shape=(120, 160, 3), drop=0.2):
def default_linear(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)
angle_out = Dense(1, activation='linear', name='angle_out')(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')
model = Model(inputs=[img_in], outputs=[angle_out], name='linear')
return model
@ -263,8 +270,7 @@ def default_categorical(input_shape=(120, 160, 3), drop=0.2):
# 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')
model = Model(inputs=[img_in], outputs=[angle_out], name='categorical')
return model
@ -278,10 +284,12 @@ if __name__ == "__main__":
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)
parser.add_argument("--model_type", type=str, default=MODEL_CATEGORICAL)
args = parser.parse_args()
params = vars(args)
train(
model_type=params["model_type"],
batch_size=params["batch_size"],
slide_size=params["slide_size"],
img_height=params["img_height"],