def train():
data = get_dataset()
x = torch.tensor(data['x'], requires_grad=True,
dtype=torch.float32) # x 实际上是位置和速度
test_x = torch.tensor(data['test_x'],
requires_grad=True,
dtype=torch.float32)
_, acce = torch.Tensor(data['test_dx']).chunk(2, 1)
_, test_acce = torch.Tensor(data['test_dx']).chunk(2, 1)
N, freedom = x.shape
freedom /= 2
input_dim = int(freedom * 2)
output_dim = int(freedom)
model_nn = MLP(input_dim, 50, output_dim, 'tanh')
model = LNN(input_dim, differentiable_model=model_nn)
optim = torch.optim.Adam(model.parameters(), 5e-3, weight_decay=1e-4)
# vanilla train loop
stats = {'train_loss': [], 'test_loss': []}
torch.autograd.set_detect_anomaly(True)
for step in range(500): #500 epoch
# train step
loss = 0
for i in range(100):
acce_hat = model.forward_new(x[i])
loss = loss + L1_loss(acce[i], acce_hat)
loss.backward()
loss /= 100
optim.step()
optim.zero_grad()
print("step {}, train_loss {:.4e}, ".format(step, loss))
writer.add_scalar('LNN/spring_train_loss', loss, step)
def train(args):
# set random seed
torch.manual_seed(args.seed)
np.random.seed(args.seed)
# init model and optimizer
if args.verbose:
print("Training baseline model:" if args.baseline else "Training HNN model:")
output_dim = args.input_dim if args.baseline else 2
nn_model = MLP(args.input_dim, args.hidden_dim, output_dim, args.nonlinearity)
model = HNN(args.input_dim, differentiable_model=nn_model,
field_type=args.field_type, baseline=args.baseline)
optim = torch.optim.Adam(model.parameters(), args.learn_rate, weight_decay=0)
# arrange data
data = get_dataset(args.name, args.save_dir, verbose=True)
x = torch.tensor( data['coords'], requires_grad=True, dtype=torch.float32)
test_x = torch.tensor( data['test_coords'], requires_grad=True, dtype=torch.float32)
dxdt = torch.Tensor(data['dcoords'])
test_dxdt = torch.Tensor(data['test_dcoords'])
# vanilla train loop
stats = {'train_loss': [], 'test_loss': []}
for step in range(args.total_steps+1):
# train step
ixs = torch.randperm(x.shape[0])[:args.batch_size]
dxdt_hat = model.time_derivative(x[ixs])
dxdt_hat += args.input_noise * torch.randn(*x[ixs].shape) # add noise, maybe
loss = L2_loss(dxdt[ixs], dxdt_hat)
loss.backward()
grad = torch.cat([p.grad.flatten() for p in model.parameters()]).clone()
optim.step() ; optim.zero_grad()
# run test data
test_ixs = torch.randperm(test_x.shape[0])[:args.batch_size]
test_dxdt_hat = model.time_derivative(test_x[test_ixs])
test_dxdt_hat += args.input_noise * torch.randn(*test_x[test_ixs].shape) # add noise, maybe
test_loss = L2_loss(test_dxdt[test_ixs], test_dxdt_hat)
# logging
stats['train_loss'].append(loss.item())
stats['test_loss'].append(test_loss.item())
if args.verbose and step % args.print_every == 0:
print("step {}, train_loss {:.4e}, test_loss {:.4e}, grad norm {:.4e}, grad std {:.4e}"
.format(step, loss.item(), test_loss.item(), grad@grad, grad.std()))
train_dxdt_hat = model.time_derivative(x)
train_dist = (dxdt - train_dxdt_hat)**2
test_dxdt_hat = model.time_derivative(test_x)
test_dist = (test_dxdt - test_dxdt_hat)**2
print('Final train loss {:.4e} +/- {:.4e}\nFinal test loss {:.4e} +/- {:.4e}'
.format(train_dist.mean().item(), train_dist.std().item()/np.sqrt(train_dist.shape[0]),
test_dist.mean().item(), test_dist.std().item()/np.sqrt(test_dist.shape[0])))
return model, stats
def get_hnn_model(args, baseline):
output_dim = args.input_dim if args.baseline else 2
nn_model = MLP(args.input_dim, 400, output_dim, args.nonlinearity)
model = HNN(args.input_dim,
differentiable_model=nn_model,
field_type=args.field_type,
baseline=args.baseline)
label = '-baseline' if args.baseline else '-hnn'
label = label + '-rad' if args.rad else label
path = '{}/{}{}.tar'.format(args.save_dir, args.name, label)
return model
def train(args):
device = torch.device(
'cuda:' + str(args.gpu) if torch.cuda.is_available() else 'cpu')
# reproducibility: set random seed
torch.manual_seed(args.seed)
np.random.seed(args.seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# init model and optimizer
if args.verbose:
print("Start training with num of points = {} and solver {}.".format(
args.num_points, args.solver))
if args.structure == False and args.baseline == True:
nn_model = MLP(args.input_dim, 600, args.input_dim, args.nonlinearity)
model = SymODEN_R(args.input_dim,
H_net=nn_model,
device=device,
baseline=True)
elif args.structure == False and args.baseline == False:
H_net = MLP(args.input_dim, 400, 1, args.nonlinearity)
g_net = MLP(int(args.input_dim / 2), 200, int(args.input_dim / 2))
model = SymODEN_R(args.input_dim,
H_net=H_net,
g_net=g_net,
device=device,
baseline=False)
elif args.structure == True and args.baseline == False:
M_net = MLP(int(args.input_dim / 2), 300, int(args.input_dim / 2))
V_net = MLP(int(args.input_dim / 2), 50, 1)
g_net = MLP(int(args.input_dim / 2), 200, int(args.input_dim / 2))
model = SymODEN_R(args.input_dim,
M_net=M_net,
V_net=V_net,
g_net=g_net,
device=device,
baseline=False,
structure=True)
else:
raise RuntimeError(
'argument *baseline* and *structure* cannot both be true')
num_parm = get_model_parm_nums(model)
print('model contains {} parameters'.format(num_parm))
optim = torch.optim.Adam(model.parameters(),
args.learn_rate,
weight_decay=1e-4)
data = get_dataset(seed=args.seed)
# modified to use the hnn stuff
x = torch.tensor(data['x'], requires_grad=True,
dtype=torch.float32) # [1125, 2] Bx2
# append zero control
u = torch.zeros_like(x[:, 0]).unsqueeze(-1)
x = torch.cat((x, u), -1)
test_x = torch.tensor(data['test_x'],
requires_grad=True,
dtype=torch.float32)
# append zero control
test_x = torch.cat((test_x, u), -1)
dxdt = torch.Tensor(data['dx']) # [1125, 2] Bx2
test_dxdt = torch.Tensor(data['test_dx'])
# training loop
stats = {'train_loss': [], 'test_loss': []}
for step in range(args.total_steps + 1):
# modified to match hnn
dq, dp, du = model.time_derivative(x).split(1, 1)
dxdt_hat = torch.cat((dq, dp), -1)
loss = L2_loss(dxdt, dxdt_hat)
loss.backward()
optim.step()
optim.zero_grad()
# run test data
dq_test, dp_test, du_test = model.time_derivative(test_x).split(1, 1)
test_dxdt_hat = torch.cat((dq_test, dp_test), -1)
test_loss = L2_loss(test_dxdt, test_dxdt_hat)
# logging
stats['train_loss'].append(loss.item())
stats['test_loss'].append(test_loss.item())
if args.verbose and step % args.print_every == 0:
print("step {}, train_loss {:.4e}, test_loss {:.4e}".format(
step, loss.item(), test_loss.item()))
train_dq, train_dp, train_du = model.time_derivative(x).split(1, 1)
train_dxdt_hat = torch.cat((train_dq, train_dp), -1)
train_dist = (dxdt - train_dxdt_hat)**2
test_dq, test_dp, test_du = model.time_derivative(test_x).split(1, 1)
test_dxdt_hat = torch.cat((test_dq, test_dp), -1)
test_dist = (test_dxdt - test_dxdt_hat)**2
print(
'Final train loss {:.4e} +/- {:.4e}\nFinal test loss {:.4e} +/- {:.4e}'
.format(train_dist.mean().item(),
train_dist.std().item() / np.sqrt(train_dist.shape[0]),
test_dist.mean().item(),
test_dist.std().item() / np.sqrt(test_dist.shape[0])))
return model, stats
def train(args):
if torch.cuda.is_available() and not args.cpu:
device = torch.device("cuda:0")
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.empty_cache()
print("Running on the GPU")
else:
device = torch.device("cpu")
print("Running on the CPU")
# set random seed
torch.manual_seed(args.seed)
np.random.seed(args.seed)
# get dataset (no test data for now)
angular_velo, acc_1, acc_2, sound = get_dataset_split(
args.folder,
args.speed,
scaled=args.scaled,
experiment_dir=args.experiment_dir,
tensor=True)
sub_col = {
0: [angular_velo, 1, 'v'],
1: [acc_1, 3, 'a1'],
2: [acc_2, 3, 'a2'],
3: [sound, 1, 's']
}
col2use = sub_col[args.sub_columns][0]
# using universal autoencoder, pre-encode the training points
autoencoder = MLPAutoencoder(sub_col[args.sub_columns][1],
args.hidden_dim,
args.latent_dim * 2,
dropout_rate=args.dropout_rate_ae)
full_model = PixelHNN(args.latent_dim * 2,
args.hidden_dim,
autoencoder=autoencoder,
nonlinearity=args.nonlinearity,
baseline=args.baseline,
dropout_rate=args.dropout_rate)
path = "{}/saved_models/{}-{}.tar".format(args.save_dir, args.ae_path,
sub_col[args.sub_columns][2])
full_model.load_state_dict(torch.load(path))
full_model.eval()
autoencoder_model = full_model.autoencoder
gcoords = autoencoder_model.encode(col2use).cpu().detach().numpy()
x = torch.tensor(gcoords, dtype=torch.float, requires_grad=True)
dx_np = full_model.time_derivative(
torch.tensor(gcoords, dtype=torch.float,
requires_grad=True)).cpu().detach().numpy()
dx = torch.tensor(dx_np, dtype=torch.float)
nnmodel = MLP(args.input_dim, args.hidden_dim, args.output_dim)
model = HNN(2, nnmodel)
model.to(device)
optim = torch.optim.Adam(model.parameters(),
args.learn_rate,
weight_decay=args.weight_decay)
print("Data from {} {}, column: {}".format(args.folder, args.speed,
sub_col[args.sub_columns][2]))
# x = torch.tensor(col2use[:-1], dtype=torch.float)
# x_next = torch.tensor(col2use[1:], dtype=torch.float)
#
# autoencoder = MLPAutoencoder(sub_col[args.sub_columns][1], args.hidden_dim, args.latent_dim * 2, dropout_rate=args.dropout_rate_ae)
# model = PixelHNN(args.latent_dim * 2, args.hidden_dim,
# autoencoder=autoencoder, nonlinearity=args.nonlinearity, baseline=args.baseline, dropout_rate=args.dropout_rate)
# model.to(device)
# optim = torch.optim.Adam(model.parameters(), args.learn_rate, weight_decay=args.weight_decay)
# vanilla ae train loop
stats = {'train_loss': []}
for step in range(args.total_steps + 1):
# train step
ixs = torch.randperm(x.shape[0])[:args.batch_size]
x_train, dxdt = x[ixs].to(device), dx[ixs].to(device)
dxdt_hat = model.time_derivative(x_train)
loss = L2_loss(dxdt, dxdt_hat)
loss.backward()
optim.step()
optim.zero_grad()
stats['train_loss'].append(loss.item())
if step % args.print_every == 0:
print("step {}, train_loss {:.4e}".format(step, loss.item()))
# train_dist = hnn_ae_loss(x, x_next, model, return_scalar=False)
# print('Final train loss {:.4e} +/- {:.4e}'
# .format(train_dist.mean().item(), train_dist.std().item() / np.sqrt(train_dist.shape[0])))
return model