diff --git a/experiment_params/train_config_default.yaml b/experiment_params/train_config_default.yaml index 903ddd5..accde2c 100644 --- a/experiment_params/train_config_default.yaml +++ b/experiment_params/train_config_default.yaml @@ -30,10 +30,6 @@ networks: n_filters: [64, 64, 64] kernel_sizes: [3, 3, 3, 3] -# Define HGN Integrator -integrator: - method: "Leapfrog" - # Define optimization optimization: epochs: 5 diff --git a/hamiltonian_generative_network.py b/hamiltonian_generative_network.py index 4637dfb..02dfde8 100644 --- a/hamiltonian_generative_network.py +++ b/hamiltonian_generative_network.py @@ -23,11 +23,11 @@ def __init__(self, transformer, hnn, decoder, - integrator, device, dtype, seq_len, - channels): + channels, + delta_t=0.125): """Instantiate a Hamiltonian Generative Network. Args: @@ -35,7 +35,6 @@ def __init__(self, transformer (networks.transformer_net.TransformerNet): Transformer neural network. hnn (networks.hamiltonian_net.HamiltonianNet): Hamiltonian neural network. decoder (networks.decoder_net.DecoderNet): Decoder neural network. - integrator (Integrator): HGN integrator. device (str): String with the device to use. E.g. 'cuda:0', 'cpu'. dtype (torch.dtype): Data type used for the networks. seq_len (int): Number of frames in each rollout. @@ -52,8 +51,8 @@ def __init__(self, self.transformer = transformer self.hnn = hnn self.decoder = decoder - self.integrator = integrator - + self.delta_t = delta_t + def forward(self, rollout_batch, n_steps=None, variational=True): """Get the prediction of the HGN for a given rollout_batch of n_steps. @@ -95,18 +94,14 @@ def forward(self, rollout_batch, n_steps=None, variational=True): # Estimate predictions for _ in range(n_steps - 1): # Compute next state - q, p = self.integrator.step(q=q, p=p, hnn=self.hnn) + delta_q, delta_p = self.hnn(q=q, p=p) + q = q + delta_q * self.delta_t + p = p + delta_p * self.delta_t prediction.append_state(q=q, p=p) - prediction.append_energy(self.integrator.energy) # This is the energy of previous timestep # Compute state reconstruction x_reconstructed = self.decoder(q) prediction.append_reconstruction(x_reconstructed) - - # We need to add the energy of the system at the last time-step - with torch.no_grad(): - last_energy = self.hnn(q=q, p=p).detach().cpu().numpy() - prediction.append_energy(last_energy) # This is the energy of previous timestep return prediction def load(self, directory): diff --git a/networks/hamiltonian_net.py b/networks/hamiltonian_net.py index e5b9009..055097c 100644 --- a/networks/hamiltonian_net.py +++ b/networks/hamiltonian_net.py @@ -64,6 +64,7 @@ def __init__(self, raise ValueError( 'Args hidden_conv_layers, n_filters, kernel_sizes, and strides' 'can only be either all None, or all defined by the user.') + self.in_shape = in_shape in_channels = in_shape[0] * 2 paddings = [int(k / 2) for k in kernel_sizes] self.in_conv = nn.Conv2d(in_channels=in_channels, @@ -83,13 +84,9 @@ def __init__(self, out_size = int((out_size - kernel_sizes[i] + 2 * paddings[i]) / strides[i]) + 1 self.out_conv = nn.Conv2d(in_channels=n_filters[-1], - out_channels=n_filters[-1], + out_channels=in_channels, kernel_size=kernel_sizes[-1], padding=paddings[-1]) - out_size = int( - (out_size - kernel_sizes[-1] + 2 * paddings[-1]) / strides[-1]) + 1 - self.n_flat = (out_size**2) * n_filters[-1] - self.linear = nn.Linear(in_features=self.n_flat, out_features=1) self.activation = act_func self.type(dtype) @@ -111,13 +108,7 @@ def forward(self, q, p): x = self.activation(self.in_conv(x)) for layer in self.hidden_layers: x = self.activation(layer(x)) - x = self.activation(self.out_conv(x)) - x = x.view(-1, self.n_flat) - x = self.linear(x) - return x - + x = self.out_conv(x) + q, p = x[:, :self.in_shape[0]], x[:, self.in_shape[0]:] + return q, p -if __name__ == '__main__': - hamiltonian_net = HamiltonianNet(in_shape=(16, 4, 4)) - q, p = torch.randn((2, 128, 16, 4, 4)) - h = hamiltonian_net(q, p) diff --git a/train.py b/train.py index 7b59135..9187fd2 100644 --- a/train.py +++ b/train.py @@ -54,7 +54,7 @@ def __init__(self, params, resume=False): # Load hgn from parameters to deice self.hgn = load_hgn(params=self.params, device=self.device, - dtype=self.dtype) + dtype=self.dtype,) # Either generate data on-the-fly or load the data from disk if "train_data" in self.params["dataset"]: @@ -162,7 +162,6 @@ def training_step(self, rollouts): train_loss.backward() self.optimizer.step() - return losses, hgn_output def fit(self): diff --git a/utilities/loader.py b/utilities/loader.py index 3f0856f..670e0a4 100644 --- a/utilities/loader.py +++ b/utilities/loader.py @@ -40,20 +40,16 @@ def load_hgn(params, device, dtype): **params["networks"]["decoder"], dtype=dtype).to(device) - # Define HGN integrator - integrator = Integrator(delta_t=params["dataset"]["rollout"]["delta_time"], - method=params["integrator"]["method"]) - # Instantiate Hamiltonian Generative Network hgn = HGN(encoder=encoder, transformer=transformer, hnn=hnn, decoder=decoder, - integrator=integrator, device=device, dtype=dtype, seq_len=params["dataset"]["rollout"]["seq_length"], - channels=params["dataset"]["rollout"]["n_channels"]) + channels=params["dataset"]["rollout"]["n_channels"], + delta_t=params["dataset"]["rollout"]["delta_time"]) return hgn diff --git a/utilities/training_logger.py b/utilities/training_logger.py index 4d8d9e4..5bfe69f 100644 --- a/utilities/training_logger.py +++ b/utilities/training_logger.py @@ -43,9 +43,6 @@ def step(self, losses, rollout_batch, prediction, model): for loss_name, loss_value in losses.items(): if loss_value is not None: self.writer.add_scalar(f'{loss_name}', loss_value, self.iteration) - enery_mean, energy_std = prediction.get_energy() - self.writer.add_scalar(f'energy/mean', enery_mean, self.iteration) - self.writer.add_scalar(f'energy/std', energy_std, self.iteration) if self.iteration % self.rollout_freq == 0: self.writer.add_video('data/input',