model.compile(loss="mean_squared_error", optimizer="rmsprop")
# print(model.summary())
# Train on events
for event_id, hits, cells, particles, truth in load_dataset(
dataset_path.get_path()):
print("-" * 15, "processing event #", event_id, "-" * 15)
transform = transformations[transformation_idx]
getattr(transforms,
transform)(hits,
True) # equivalent to transforms.transform(hits, True)
true_tracks, hit_tracks, max_len = formatter.getSortedTracks(
particles, truth, hits)
x, y = formatter.getInputOutput(hit_tracks, 1, 3, seq_len, seq_len, 18)
## Train the LSTM model ##
print("Training with the", transformations[transformation_idx],
"tranformation")
model.fit(x,
y,
batch_size=batch_size,
epochs=num_epochs,
validation_split=0)
model.save(transformations[transformation_idx] + str(seq_len) + ".keras")
## debugging code ##
# sanity checks for shape
# printShapes = False
