NASA RUL project - lstm training

train_lstm.py

checkpoint_callback = ModelCheckpoint(
    dirpath='checkpoints',
    filename='best-checkpoint',
    save_top_k=1,
    verbose=True,
    monitor='val_loss',
    mode='min'
)

logger = TensorBoardLogger('lightning_logs', name='RUL')
trainer = pl.Trainer(
    logger=logger,
    callbacks=[checkpoint_callback],
    max_epochs=n_epochs,
    gpus=1,
    progress_bar_refresh_rate=30
)

trainer.fit(model, data_module)

trained_model = RULPredictor.load_from_checkpoint(
    trainer.checkpoint_callback.best_model_path,
    n_features=len(feature_columns)
)
trained_model.freeze()