Torch layer

Torch layer#

The lowest layer is the torch layer. This consists of the torch dataset, dataloader, optimizers, and schedulers.

Datasets#

In the physicsml package most models are graph based, meaning that they can all share the same low-level graph dataset. Once the featurised data is given to the train or predict methods, it is internally converted to a torch compatible dataset. This dataset is built on top of torch-geometric (see their docs), an excellent library for graph-based torch modelling. In general, a model developer will rarely have to worry about this (unless they are implementing a special kind of model that requires extra inputs).

The only exception to this is the ANI models. They require a special dataset (which is simpler) since the only inputs they use are species and coordinates.

Dataloaders#

Again, we use the torch-geometric dataloading for all models except for the ANI models. Model developers will rarely have to modify this unless their models require additional inputs. Here, we will describe what a batch from each dataloader looks like and what you need to expect as an input to your model.

torch-geometric batch#

A torch-geometric batch contains the following key-value pairs

  • num_nodes: Shape = torch.Size([]), Type = torch.int64

    The total number of nodes in the batch.

  • edge_index: Shape = torch.Size([2, num_edges]) Type = torch.int64.

    The edge indices (receiver, sender). They are concatenated with a cumulative shift of the number of nodes in previous graphs (for batch GNN operations like scatter).

  • node_attrs: Shape = torch.Size([num_nodes, dim_node_attrs]) Type = torch.float32

    The node attributes (i.e. initial node features). It is a good idea to keep these distinct from node features (for downstream operations which require the original node attributes).

  • edge_attrs: Shape = torch.Size([num_edges, dim_edge_attrs]) Type = torch.float32

    The edge attributes (i.e. initial edge features). It is a good idea to keep these distinct from edge features (for downstream operations which require the original edge attributes).

  • graph_attrs: Shape = torch.Size([num_graphs, dim_graph_attrs]) Type = torch.float32

    The graph attributes.

  • coordinates: Shape = torch.Size([num_nodes, 3]) Type = torch.float32

    The coordinates.

  • total_atomic_energy: Shape = torch.Size([num_graphs]) Type = torch.float32

    The summed up atomic energies.

  • raw_atomic_numbers: Shape = torch.Size([num_nodes]) Type = torch.int64

    The raw mapped atomic numbers.

  • atomic_numbers: Shape = torch.Size([num_nodes, num_elements]) Type = torch.float32

    The one-hot encoded atomic numbers.

  • cell: Shape = torch.Size([3, 3]) Type = torch.float32

    The unit cell dimensions. This is used to compute distance and vectors when working in periodic systems.

  • cell_shift_vector: Shape = torch.Size([num_edges, 3]) Type = torch.float32

    The shift vectors for the connected edges used to compute distances and vectors when working in periodic systems.

  • y_graph_scalars: Shape = torch.Size([num_graphs, y_graph_scalars_num_tasks]) Type = torch.float32

    The graph level scalars references for the predictions.

  • y_edge_scalars: Shape = torch.Size([num_edges, y_edge_scalars_num_tasks]) Type = torch.float32

    The edge level scalars references for the predictions.

  • y_node_scalars: Shape = torch.Size([num_nodes, y_node_scalars_num_tasks]) Type = torch.float32

    The node level scalars references for the predictions.

  • y_graph_vector: Shape = torch.Size([num_graphs, 3]) Type = torch.float32

    The graph level vector references for the predictions.

  • y_edge_vector: Shape = torch.Size([num_edge, 3]) Type = torch.float32

    The edge level vector references for the predictions.

  • y_node_vector: Shape = torch.Size([num_nodes, 3]) Type = torch.float32

    The node level vector references for the predictions.

  • num_graphs: Shape = torch.Size([]) Type = torch.int64

    The number of graphs in the batch.

  • batch: Shape = torch.Size([num_nodes]) Type = torch.int64

    The batch index of each node (which graph it belongs to).

  • ptr: Shape = torch.Size([num_graphs + 1]) Type = torch.int64

    The cumulative number of nodes in each graph (with an extra 0 at the beginning).

An ANI model batch looks like

  • species: Shape = torch.Size([num_graphs, max_num_nodes]) Type = torch.int64

    The mapped atom types padded up to the max number of nodes in the batch.

  • coordinates: Shape = torch.Size([num_graphs, max_num_nodes, 3]) Type = torch.float32

    The coordinates.

  • total_atomic_energy: Shape = torch.Size([num_graphs]) Type = torch.float32

    The summed up atomic energies.

  • y_graph_scalars: Shape = torch.Size([num_graphs, y_graph_scalars_num_tasks]) Type = torch.float32

    The graph level scalars references for the predictions.

  • y_node_vector: Shape = torch.Size([num_nodes, 3]) Type = torch.float32

    The node level vector references for the predictions.

By default the precision for floating point number is float32 unless otherwise specified in the Trainer config

optimizer#

All optimizers from torch.optim can be used. We have created a serialisation for loading them via configs. The config has the following structure

  • name: The name of optimizer (from the torch.optim module).

  • config: A dict of the kwargs the optimizer takes.

scheduler config#

All schedulers from torch.optim.lr_scheduler can be used. We have created a serialisation for loading them via configs. The config has the following structure

  • name: The name of scheduler (from the torch.optim.lr_scheduler module).

  • config: A dict of the kwargs the scheduler takes.

  • interval: The interval to use (for schedulers like ReduceLROnPlateua).

  • frequency: The interval frequency to use (for schedulers like ReduceLROnPlateua).

  • monitor: The metric to monitor (for schedulers like ReduceLROnPlateua).

  • strict: Whether to fail or raise a warning if monitoring metric not found (for schedulers like ReduceLROnPlateua).

Contents