if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
return args
if __name__ == '__main__':
args = get_args()
hps = get_default_hparams()
with open("./hps/hps.pickle", "wb") as output_file:
pickle.dump(hps, output_file)
M = ModelDesc(hps)
logger.auto_set_dir(action='d')
ds_train, ds_test = get_data(hps)
Trainer(input=QueueInput(ds_train), model=M).train_with_defaults(
callbacks=[
ModelSaver(),
callbacks.MergeAllSummaries(),
MinSaver('total_loss'),
InferenceRunner(ds_test, [
ScalarStats('predict_trend/mse_predict_loss'),
])
],
steps_per_epoch=hps.steps_per_epoch,
max_epoch=hps.epochs,
session_init=SaverRestore(args.load) if args.load else None)
if __name__ == '__main__':
args = get_args()
hps = get_default_hparams()
with open("./hps/hps.pickle", "wb") as output_file:
pickle.dump(hps, output_file)
M = ModelDesc(hps)
logger.auto_set_dir(action='d')
ds_train, ds_test = get_data(hps)
Trainer(input=QueueInput(ds_train), model=M).train_with_defaults(
callbacks=[
ModelSaver(),
callbacks.MergeAllSummaries(),
MinSaver('total_loss'),
# InferenceRunner(ds_test, [ScalarStats('predict_trend/accuracy_')])
InferenceRunner(ds_test, [
ScalarStats('predict_trend/accuracy_'),
BinaryClassificationStats(
pred_tensor_name='predict_trend/y_pred_one_hot',
label_tensor_name='y_one_hot')
])
],
steps_per_epoch=hps.steps_per_epoch,
max_epoch=hps.epochs,
session_init=SaverRestore(args.load) if args.load else None)