def main(): __model = MLP(in_dim=400, hidden_dim=HIDDEN_DIM, out_dim=800) print("\tTraining {}...".format(__model.__str__().split("(")[0])) optim = torch.optim.Adam(__model.parameters(), lr=LEARNING_RATE) train(__model, optim, X_train, y_train, 128) mse = test(__model, X_test, y_test).mean().item() print(mse)
def main(): save_dir = './results/' models = [ MLP( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1, activation='relu6', ), MLP( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1, activation='none', ), NAC( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1, ), NALU( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1 ), ] results = {} for fn_str, fn in ARITHMETIC_FUNCTIONS.items(): print('[*] Testing function: {}'.format(fn_str)) results[fn_str] = [] # dataset X_train, y_train, X_test, y_test = generate_data( num_train=500, num_test=50, dim=100, num_sum=5, fn=fn, support=RANGE, ) # random model random_mse = [] for i in range(100): net = MLP( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1, activation='relu6', ) mse = test(net, X_test, y_test) random_mse.append(mse.mean().item()) results[fn_str].append(np.mean(random_mse)) # others for net in models: print("\tTraining {}...".format(net.__str__().split("(")[0])) optim = torch.optim.RMSprop(net.parameters(), lr=LEARNING_RATE) train(net, optim, X_train, y_train, NUM_ITERS) mse = test(net, X_test, y_test).mean().item() results[fn_str].append(mse) with open(save_dir + "interpolation.txt", "w") as f: f.write("Relu6\tNone\tNAC\tNALU\n") for k, v in results.items(): rand = results[k][0] mses = [100.0*x/rand for x in results[k][1:]] if NORMALIZE: f.write("{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}\n".format(*mses)) else: f.write("{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}\n".format(*results[k][1:]))