Пример #1
0
    def __init__(self,
                 dim,
                 hidden_dim,
                 autoencoder,
                 nonlinearity='tanh',
                 dt=1e-3,
                 device=None):
        super(PixelLagrangian, self).__init__()
        self.autoencoder = autoencoder

        S_net = MLP(dim, hidden_dim, dim**2, nonlinearity).to(device)
        U_net = MLP(dim, hidden_dim, 1, nonlinearity).to(device)
        self.lag = LagrangianFriction(dim, S_net, U_net, dt=dt, device=device)
Пример #2
0
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)
Пример #3
0
    def __init__(self, input_dim, hidden_dim, autoencoder,
                 field_type='solenoidal', nonlinearity='tanh', baseline=False):
        super(PixelHNN, self).__init__()
        self.autoencoder = autoencoder
        self.baseline = baseline

        output_dim = input_dim if baseline else 2
        nn_model = MLP(input_dim, hidden_dim, output_dim, nonlinearity)
        self.hnn = HNN(input_dim, differentiable_model=nn_model, field_type=field_type, baseline=baseline)
Пример #4
0
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
Пример #5
0
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
Пример #6
0
    def __init__(self,
                 dim,
                 autoencoder,
                 nonlinearity='tanh',
                 dt=1e-3,
                 M_hidden=300,
                 V_hidden=50,
                 g_hidden=200,
                 device=None):
        super(PixelSymODEN_R, self).__init__()
        self.autoencoder = autoencoder

        M_net = PSD(dim, M_hidden, dim).to(device)
        V_net = MLP(dim, V_hidden, 1).to(device)
        g_net = MLP(dim, g_hidden, dim).to(device)
        self.symoden = SymODEN_R(dim * 2,
                                 M_net=M_net,
                                 V_net=V_net,
                                 g_net=g_net,
                                 device=device,
                                 baseline=False,
                                 structure=True)
Пример #7
0
def get_model(args, baseline, structure, naive, num_points):
    M_net = PSD(3, 400, 2).to(device)
    g_net = MatrixNet(3, 300, 4, shape=(2,2)).to(device)
    if structure == False:
        if naive and baseline:
            raise RuntimeError('argument *baseline* and *naive* cannot both be true')
        elif naive:
            input_dim = 6
            output_dim = 5
            nn_model = MLP(input_dim, 1000, output_dim, args.nonlinearity).to(device)
            model = SymODEN_R1_T1(args.num_angle, H_net=nn_model, device=device, baseline=baseline, naive=naive, u_dim=2)
        elif baseline:
            input_dim = 6
            output_dim = 4
            nn_model = MLP(input_dim, 700, output_dim, args.nonlinearity).to(device)
            model = SymODEN_R1_T1(args.num_angle, H_net=nn_model, M_net=M_net, device=device, baseline=baseline, naive=naive, u_dim=2)
        else:
            input_dim = 5
            output_dim = 1
            nn_model = MLP(input_dim, 500, output_dim, args.nonlinearity).to(device)
            model = SymODEN_R1_T1(args.num_angle, H_net=nn_model, M_net=M_net, g_net=g_net, device=device, baseline=baseline, naive=naive, u_dim=2)
    elif structure == True and baseline ==False and naive==False:
        V_net = MLP(3, 300, 1).to(device)
        model = SymODEN_R1_T1(args.num_angle, M_net=M_net, V_net=V_net, g_net=g_net, device=device, baseline=baseline, structure=True, u_dim=2).to(device)
    else:
        raise RuntimeError('argument *structure* is set to true, no *baseline* or *naive*!')

    if naive:
        label = '-naive_ode'
    elif baseline:
        label = '-baseline_ode'
    else:
        label = '-hnn_ode'
    struct = '-struct' if structure else ''
    path = '{}/{}{}{}-{}-p{}.tar'.format(args.save_dir, args.name, label, struct, args.solver, args.num_points)
    model.load_state_dict(torch.load(path, map_location=device))
    path = '{}/{}{}{}-{}-p{}-stats.pkl'.format(args.save_dir, args.name, label, struct, args.solver, args.num_points)
    stats = from_pickle(path)
    return model, stats
Пример #8
0
def get_model(args, baseline, structure, damping, num_points, gym=False):
    if structure == False and baseline == True:
        nn_model = MLP(args.input_dim, 600, args.input_dim,
                       args.nonlinearity).to(device)
        model = SymODEN_R(args.input_dim,
                          H_net=nn_model,
                          device=device,
                          baseline=True)
    elif structure == False and baseline == False:
        H_net = MLP(args.input_dim, 400, 1, args.nonlinearity).to(device)
        g_net = MLP(int(args.input_dim / 2), 200,
                    int(args.input_dim / 2)).to(device)
        model = SymODEN_R(args.input_dim,
                          H_net=H_net,
                          g_net=g_net,
                          device=device,
                          baseline=False)
    elif structure == True and baseline == False:
        # M_net = MLP(1, args.hidden_dim, 1).to(device)
        M_net = MLP(int(args.input_dim / 2), 300, int(args.input_dim / 2))
        V_net = MLP(int(args.input_dim / 2), 50, 1).to(device)
        g_net = MLP(int(args.input_dim / 2), 200,
                    int(args.input_dim / 2)).to(device)
        model = SymODEN_R(args.input_dim,
                          M_net=M_net,
                          V_net=V_net,
                          g_net=g_net,
                          device=device,
                          baseline=False,
                          structure=True).to(device)
    else:
        raise RuntimeError(
            'argument *baseline* and *structure* cannot both be true')
    model_name = 'baseline_ode' if baseline else 'hnn_ode'
    struct = '-struct' if structure else ''
    rad = '-rad' if args.rad else ''
    path = '{}pend-{}{}-{}-p{}{}.tar'.format(args.save_dir, model_name, struct,
                                             args.solver, num_points, rad)
    model.load_state_dict(torch.load(path, map_location=device))
    path = '{}/pend-{}{}-{}-p{}-stats{}.pkl'.format(args.save_dir, model_name,
                                                    struct, args.solver,
                                                    num_points, rad)
    stats = from_pickle(path)
    return model, stats
Пример #9
0
def train(args):
    # import ODENet
    # from torchdiffeq import odeint
    from torchdiffeq import odeint_adjoint as odeint

    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).to(device)    
        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).to(device)
        g_net = MLP(int(args.input_dim/2), 200, int(args.input_dim/2)).to(device)
        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).to(device)
        g_net = MLP(int(args.input_dim/2), 200, int(args.input_dim/2)).to(device)
        model = SymODEN_R(args.input_dim, M_net=M_net, V_net=V_net, g_net=g_net, device=device, baseline=False, structure=True).to(device)
    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)

    # arrange data
    us = [-2.0, -1.0, 0.0, 1.0, 2.0]
    # us = [0.0]
    data = get_dataset(seed=args.seed, timesteps=45,
                    save_dir=args.save_dir, rad=args.rad, us=us, samples=50) 
    train_x, t_eval = arrange_data(data['x'], data['t'], num_points=args.num_points)
    test_x, t_eval = arrange_data(data['test_x'], data['t'], num_points=args.num_points)

    train_x = torch.tensor(train_x, requires_grad=True, dtype=torch.float32).to(device) 
    test_x = torch.tensor(test_x, requires_grad=True, dtype=torch.float32).to(device)
    t_eval = torch.tensor(t_eval, requires_grad=True, dtype=torch.float32).to(device)

    # training loop
    stats = {'train_loss': [], 'test_loss': [], 'forward_time': [], 'backward_time': [],'nfe': []}
    for step in range(args.total_steps+1):
        train_loss = 0
        test_loss = 0
        for i in range(train_x.shape[0]):
            
            t = time.time()
            train_x_hat = odeint(model, train_x[i, 0, :, :], t_eval, method=args.solver)            
            forward_time = time.time() - t
            train_loss_mini = L2_loss(train_x[i,:,:,:], train_x_hat)
            train_loss = train_loss + train_loss_mini

            t = time.time()
            train_loss_mini.backward() ; optim.step() ; optim.zero_grad()
            backward_time = time.time() - t

            # run test data
            test_x_hat = odeint(model, test_x[i, 0, :, :], t_eval, method=args.solver)
            test_loss_mini = L2_loss(test_x[i,:,:,:], test_x_hat)
            test_loss = test_loss + test_loss_mini

        # logging
        stats['train_loss'].append(train_loss.item())
        stats['test_loss'].append(test_loss.item())
        stats['forward_time'].append(forward_time)
        stats['backward_time'].append(backward_time)
        stats['nfe'].append(model.nfe)
        if args.verbose and step % args.print_every == 0:
            print("step {}, train_loss {:.4e}, test_loss {:.4e}".format(step, train_loss.item(), test_loss.item()))

    # calculate loss mean and std for each traj.
    train_x, t_eval = data['x'], data['t']
    test_x, t_eval = data['test_x'], data['t']

    train_x = torch.tensor(train_x, requires_grad=True, dtype=torch.float32).to(device)
    test_x = torch.tensor(test_x, requires_grad=True, dtype=torch.float32).to(device)
    t_eval = torch.tensor(t_eval, requires_grad=True, dtype=torch.float32).to(device)

    train_loss = []
    test_loss = []
    for i in range(train_x.shape[0]):
        train_x_hat = odeint(model, train_x[i, 0, :, :], t_eval, method=args.solver)            
        train_loss.append((train_x[i,:,:,:] - train_x_hat)**2)

        # run test data
        test_x_hat = odeint(model, test_x[i, 0, :, :], t_eval, method=args.solver)
        test_loss.append((test_x[i,:,:,:] - test_x_hat)**2)

    train_loss = torch.cat(train_loss, dim=1)
    train_loss_per_traj = torch.sum(train_loss, dim=(0,2))

    test_loss = torch.cat(test_loss, dim=1)
    test_loss_per_traj = torch.sum(test_loss, dim=(0,2))

    print('Final trajectory train loss {:.4e} +/- {:.4e}\nFinal trajectory test loss {:.4e} +/- {:.4e}'
    .format(train_loss_per_traj.mean().item(), train_loss_per_traj.std().item(),
            test_loss_per_traj.mean().item(), test_loss_per_traj.std().item()))

    stats['traj_train_loss'] = train_loss_per_traj.detach().cpu().numpy()
    stats['traj_test_loss'] = test_loss_per_traj.detach().cpu().numpy()

    return model, stats
Пример #10
0
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, 400, 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=1e-4)

    # the data API is different
    # make sure it is a fair comparison
    # generate the data the same way as in the SymODEN
    # compute the time derivative based on the generated data
    us = [0.0]
    data = get_dataset(seed=args.seed, save_dir=args.save_dir, 
                    rad=args.rad, us=us, samples=50, timesteps=45) 
    
    # arrange data 

    train_x, t_eval = data['x'][0,:,:,0:2], data['t']
    test_x, t_eval = data['test_x'][0,:,:,0:2], data['t']

    train_dxdt = (train_x[1:,:,:] - train_x[:-1,:,:]) / (t_eval[1] - t_eval[0])
    test_dxdt = (test_x[1:,:,:] - test_x[:-1,:,:]) / (t_eval[1] - t_eval[0])

    train_x = train_x[0:-1,:,:].reshape((-1,2))
    test_x = test_x[0:-1,:,:].reshape((-1,2))
    test_dxdt = test_dxdt.reshape((-1,2))
    train_dxdt = train_dxdt.reshape((-1,2))
    
    x = torch.tensor( train_x, requires_grad=True, dtype=torch.float32)
    test_x = torch.tensor( test_x, requires_grad=True, dtype=torch.float32)
    dxdt = torch.Tensor(train_dxdt)
    test_dxdt = torch.Tensor(test_dxdt)

    # vanilla train loop
    stats = {'train_loss': [], 'test_loss': []}
    for step in range(args.total_steps+1):
        
        # train step
        dxdt_hat = model.rk4_time_derivative(x) if args.use_rk4 else model.time_derivative(x)
        loss = L2_loss(dxdt, dxdt_hat)
        loss.backward() ; optim.step() ; optim.zero_grad()
        
        # run test data
        test_dxdt_hat = model.rk4_time_derivative(test_x) if args.use_rk4 else model.time_derivative(test_x)
        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_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
Пример #11
0
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=1e-4)

    # arrange data
    data = get_dataset(seed=args.seed)
    x = torch.tensor(data['x'], requires_grad=True, dtype=torch.float32)
    test_x = torch.tensor(data['test_x'],
                          requires_grad=True,
                          dtype=torch.float32)
    dxdt = torch.Tensor(data['dx'])
    test_dxdt = torch.Tensor(data['test_dx'])

    # vanilla train loop
    stats = {'train_loss': [], 'test_loss': []}
    for step in range(args.total_steps + 1):

        # train step
        dxdt_hat = model.rk4_time_derivative(
            x) if args.use_rk4 else model.time_derivative(x)
        loss = L2_loss(dxdt, dxdt_hat)
        loss.backward()
        optim.step()
        optim.zero_grad()

        # run test data
        test_dxdt_hat = model.rk4_time_derivative(
            test_x) if args.use_rk4 else model.time_derivative(test_x)
        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_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
Пример #12
0
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)

    print("{} {}".format(args.folder, args.speed))
    print("Training scaled model:" if args.scaled else "Training noisy model:")
    print('{} pairs of coords in latent space '.format(args.latent_dim))

    #using universal autoencoder, pre-encode the training points
    autoencoder = MLPAutoencoder(args.input_dim_ae,
                                 args.hidden_dim,
                                 args.latent_dim * 2,
                                 nonlinearity='relu')
    full_model = PixelHNN(args.latent_dim * 2,
                          args.hidden_dim,
                          autoencoder=autoencoder,
                          nonlinearity=args.nonlinearity,
                          baseline=args.baseline)
    path = "{}/saved_models/{}.tar".format(args.save_dir, args.ae_path)
    full_model.load_state_dict(torch.load(path))
    full_model.eval()
    autoencoder_model = full_model.autoencoder

    # get dataset (no test data for now)
    data = get_dataset(args.folder,
                       args.speed,
                       scaled=args.scaled,
                       split=args.split_data,
                       experiment_dir=args.experiment_dir,
                       tensor=True)
    gcoords = autoencoder_model.encode(data).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)

    # vanilla ae 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]
        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
Пример #13
0
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:")
  S_net = MLP(int(args.input_dim/2), 140, int(args.input_dim/2)**2, args.nonlinearity)
  U_net = MLP(int(args.input_dim/2), 140, 1, args.nonlinearity)
  model = Lagrangian(int(args.input_dim/2), S_net, U_net, dt=1e-3)

  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)

  # arrange data
  data = get_lag_dataset(seed=args.seed)
  x = torch.tensor( data['x'], requires_grad=False, dtype=torch.float32)
  # 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=False, dtype=torch.float32)
  # append zero control
  test_x = torch.cat((test_x, u), -1)

  dxdt = torch.Tensor(data['dx'])
  test_dxdt = torch.Tensor(data['test_dx'])

  # vanilla train loop
  stats = {'train_loss': [], 'test_loss': []}
  for step in range(args.total_steps+1):
    
    # train step
    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
Пример #14
0
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
Пример #15
0
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