def train_test_split(self) -> Tuple[List[TubRecord], List[TubRecord]]:
print(f'Loading tubs from paths {self.tub_paths}')
self.records.clear()
img1 = None
img2 = None
img3 = None
img4 = None
for tub in self.tubs:
for underlying in tub:
if img3 == None:
img1 = underlying
img2 = underlying
img3 = underlying
img4 = underlying
else:
img1 = img2
img2 = img3
img3 = img4
img4 = underlying
record = TubRecord(self.config,
tub.base_path,
underlyings=[img1, img2, img3, img4])
self.records.append(record)
return train_test_split(self.records,
shuffle=self.shuffle,
test_size=(1. - self.config.TRAIN_TEST_SPLIT))
def setup(self, stage=None):
"""Load all the tub data and set up the datasets.
Args:
stage ([string], optional): setup expects a string arg stage.
It is used to separate setup logic for trainer.fit
and trainer.test. Defaults to None.
"""
# Loop through all the different tubs and load all the records for each of them
for tub in self.tubs:
for underlying in tub:
record = TubRecord(self.config,
tub.base_path,
underlying=underlying)
self.records.append(record)
train_records, val_records = train_test_split(
self.records, test_size=(1. - self.config.TRAIN_TEST_SPLIT))
assert len(
val_records) > 0, "Not enough validation data. Add more data"
self.train_dataset = TorchTubDataset(self.config,
train_records,
transform=self.transform)
self.val_dataset = TorchTubDataset(self.config,
val_records,
transform=self.transform)
def train_test_split(self):
print('Loading tubs from paths %s' % (self.tub_paths))
for tub in self.tubs:
for record in tub:
record['_image_base_path'] = tub.images_base_path
self.records.append(record)
return train_test_split(self.records,
shuffle=self.shuffle,
test_size=self.test_size)
def test_training_pipeline(config: Config, model_type: str,
train_filter: Callable[[TubRecord], bool]) -> None:
"""
Testing consistency of the model interfaces and data used in training
pipeline.
:param config: donkey config
:param model_type: test specification of model type
:param train_filter: filter for records
:return: None
"""
kl = get_model_by_type(model_type, config)
tub_dir = config.DATA_PATH_ALL if model_type in full_tub else \
config.DATA_PATH
# don't shuffle so we can identify data for testing
config.TRAIN_FILTER = train_filter
dataset = TubDataset(config, [tub_dir], seq_size=kl.seq_size())
training_records, validation_records = \
train_test_split(dataset.get_records(), shuffle=False,
test_size=(1. - config.TRAIN_TEST_SPLIT))
seq = BatchSequence(kl, config, training_records, True)
data_train = seq.create_tf_data()
num_whole_batches = len(training_records) // config.BATCH_SIZE
# this takes all batches into one list
tf_batch = list(data_train.take(num_whole_batches).as_numpy_iterator())
it = iter(training_records)
for xy_batch in tf_batch:
# extract x and y values from records, asymmetric in x and y b/c x
# requires image manipulations
batch_records = [next(it) for _ in range(config.BATCH_SIZE)]
records_x = [
kl.x_translate(kl.x_transform_and_process(r, normalize_image))
for r in batch_records
]
records_y = [kl.y_translate(kl.y_transform(r)) for r in batch_records]
# from here all checks are symmetrical between x and y
for batch, o_type, records \
in zip(xy_batch, kl.output_types(), (records_x, records_y)):
# check batch dictionary have expected keys
assert batch.keys() == o_type.keys(), \
'batch keys need to match models output types'
# convert record values into arrays of batch size
values = defaultdict(list)
for r in records:
for k, v in r.items():
values[k].append(v)
# now convert arrays of floats or numpy arrays into numpy arrays
np_dict = dict()
for k, v in values.items():
np_dict[k] = np.array(v)
# compare record values with values from tf.data
for k, v in batch.items():
assert np.isclose(v, np_dict[k]).all()
def train_test_split(self) -> Tuple[List[TubRecord], List[TubRecord]]:
print(f'Loading tubs from paths {self.tub_paths}')
self.records.clear()
for tub in self.tubs:
for underlying in tub:
record = TubRecord(self.config,
tub.base_path,
underlying=underlying)
self.records.append(record)
return train_test_split(self.records,
shuffle=self.shuffle,
test_size=(1. - self.config.TRAIN_TEST_SPLIT))
def train_test_split(self) -> Tuple[List[TubRecord], List[TubRecord]]:
msg = f'Loading tubs from paths {self.tub_paths}' + f' with filter ' \
f'{self.train_filter}' if self.train_filter else ''
print(msg)
self.records.clear()
for tub in self.tubs:
for underlying in tub:
record = TubRecord(self.config, tub.base_path, underlying)
if not self.train_filter or self.train_filter(record):
self.records.append(record)
return train_test_split(self.records,
shuffle=self.shuffle,
test_size=(1. - self.config.TRAIN_TEST_SPLIT))
def train(cfg: Config, tub_paths: str, model: str = None,
model_type: str = None, transfer: str = None, comment: str = None) \
-> tf.keras.callbacks.History:
"""
Train the model
"""
database = PilotDatabase(cfg)
if model_type is None:
model_type = cfg.DEFAULT_MODEL_TYPE
model_path, model_num = \
get_model_train_details(database, model)
base_path = os.path.splitext(model_path)[0]
kl = get_model_by_type(model_type, cfg)
if transfer:
kl.load(transfer)
if cfg.PRINT_MODEL_SUMMARY:
print(kl.interpreter.model.summary())
tubs = tub_paths.split(',')
all_tub_paths = [os.path.expanduser(tub) for tub in tubs]
dataset = TubDataset(config=cfg,
tub_paths=all_tub_paths,
seq_size=kl.seq_size())
training_records, validation_records \
= train_test_split(dataset.get_records(), shuffle=True,
test_size=(1. - cfg.TRAIN_TEST_SPLIT))
print(f'Records # Training {len(training_records)}')
print(f'Records # Validation {len(validation_records)}')
# We need augmentation in validation when using crop / trapeze
training_pipe = BatchSequence(kl, cfg, training_records, is_train=True)
validation_pipe = BatchSequence(kl,
cfg,
validation_records,
is_train=False)
tune = tf.data.experimental.AUTOTUNE
dataset_train = training_pipe.create_tf_data().prefetch(tune)
dataset_validate = validation_pipe.create_tf_data().prefetch(tune)
train_size = len(training_pipe)
val_size = len(validation_pipe)
### training/validation length limit. Large validation datasets cause memory leaks.
train_limit = cfg.TRAIN_LIMIT
train_len = len(training_records)
if train_limit is not None and train_len > train_limit:
train_decrease = train_limit / train_len
_train_size = math.ceil(train_size * train_decrease)
print(f'train steps decrease from {train_size} to {_train_size}')
train_size = _train_size
val_limit = cfg.VALIDATION_LIMIT
val_len = len(validation_records)
if val_limit is not None and val_len > val_limit:
val_decrease = val_limit / val_len
_val_size = math.ceil(val_size * val_decrease)
print(f'val steps decrease from {val_size} to {_val_size}')
val_size = _val_size
assert val_size > 0, "Not enough validation data, decrease the batch " \
"size or add more data."
history = kl.train(model_path=model_path,
train_data=dataset_train,
train_steps=train_size,
batch_size=cfg.BATCH_SIZE,
validation_data=dataset_validate,
validation_steps=val_size,
epochs=cfg.MAX_EPOCHS,
verbose=cfg.VERBOSE_TRAIN,
min_delta=cfg.MIN_DELTA,
use_early_stop=cfg.USE_EARLY_STOP,
patience=cfg.EARLY_STOP_PATIENCE,
show_plot=cfg.SHOW_PLOT)
if getattr(cfg, 'CREATE_TF_LITE', True):
tf_lite_model_path = f'{base_path}.tflite'
keras_model_to_tflite(model_path, tf_lite_model_path)
if getattr(cfg, 'CREATE_TENSOR_RT', False):
# load h5 (ie. keras) model
model_rt = load_model(model_path)
# save in tensorflow savedmodel format (i.e. directory)
model_rt.save(f'{base_path}.savedmodel')
# pass savedmodel to the rt converter
saved_model_to_tensor_rt(f'{base_path}.savedmodel', f'{base_path}.trt')
database_entry = {
'Number': model_num,
'Name': os.path.basename(base_path),
'Type': str(kl),
'Tubs': tub_paths,
'Time': time(),
'History': history.history,
'Transfer': os.path.basename(transfer) if transfer else None,
'Comment': comment,
'Config': str(cfg)
}
database.add_entry(database_entry)
database.write()
return history