def main(args):
train, test = read_data(args.in_dir)
train, val = split_data(train)
print('\n\nTrain: ', train.shape)
print('Val: ', val.shape)
print('Test: ', test.shape)
# RBM object
rbm = RBM(args.num_hidden, val.shape[1], args.lr, args.n, args.batch_size,
args.epochs)
# Train RBM
train_loss, val_loss = rbm.Train(train, val)
# Create output dir if it doesn't exist
if not os.path.exists(args.out_dir):
os.makedirs(args.out_dir)
# Plot error
plot_error(train_loss, val_loss, args.out_dir)
# Performance on Test set
error_test = rbm.reconstruction_error(test.T)
print("\n\n\nReconstruction error...\n")
print("Train : ", train_loss[-1])
print("Val : ", val_loss[-1])
print("Test : ", error_test)
# For viewing reconstruction
reconstruct_images(rbm, test, (args.image_height, args.image_width),
args.out_dir)
# Saving the model learned weights
pickle.dump([rbm.W, rbm.b_h, rbm.b_v],
open(args.out_dir + '\\weights.pkl', 'wb'))
print(f"\n\nRBM weights saved in {args.out_dir}\\weights.pkl")
