Neural Network Utility Functions

chemprop.nn_utils.py contains utility funtions specific to neural networks.

class chemprop.nn_utils.NoamLR(optimizer: torch.optim.optimizer.Optimizer, warmup_epochs: List[Union[float, int]], total_epochs: List[int], steps_per_epoch: int, init_lr: List[float], max_lr: List[float], final_lr: List[float])[source]

Noam learning rate scheduler with piecewise linear increase and exponential decay.

The learning rate increases linearly from init_lr to max_lr over the course of the first warmup_steps (where warmup_steps = warmup_epochs * steps_per_epoch). Then the learning rate decreases exponentially from max_lr to final_lr over the course of the remaining total_steps - warmup_steps (where total_steps = total_epochs * steps_per_epoch). This is roughly based on the learning rate schedule from Attention is All You Need, section 5.3.

Parameters
  • optimizer – A PyTorch optimizer.

  • warmup_epochs – The number of epochs during which to linearly increase the learning rate.

  • total_epochs – The total number of epochs.

  • steps_per_epoch – The number of steps (batches) per epoch.

  • init_lr – The initial learning rate.

  • max_lr – The maximum learning rate (achieved after warmup_epochs).

  • final_lr – The final learning rate (achieved after total_epochs).

get_lr() List[float][source]

Gets a list of the current learning rates.

Returns

A list of the current learning rates.

step(current_step: Optional[int] = None)[source]

Updates the learning rate by taking a step.

Parameters

current_step – Optionally specify what step to set the learning rate to. If None, current_step = self.current_step + 1.

chemprop.nn_utils.compute_gnorm(model: torch.nn.modules.module.Module) float[source]

Computes the norm of the gradients of a model.

Parameters

model – A PyTorch model.

Returns

The norm of the gradients of the model.

chemprop.nn_utils.compute_pnorm(model: torch.nn.modules.module.Module) float[source]

Computes the norm of the parameters of a model.

Parameters

model – A PyTorch model.

Returns

The norm of the parameters of the model.

chemprop.nn_utils.get_activation_function(activation: str) torch.nn.modules.module.Module[source]

Gets an activation function module given the name of the activation.

Supports:

  • ReLU

  • LeakyReLU

  • PReLU

  • tanh

  • SELU

  • ELU

Parameters

activation – The name of the activation function.

Returns

The activation function module.

chemprop.nn_utils.index_select_ND(source: torch.Tensor, index: torch.Tensor) torch.Tensor[source]

Selects the message features from source corresponding to the atom or bond indices in index.

Parameters
  • source – A tensor of shape (num_bonds, hidden_size) containing message features.

  • index – A tensor of shape (num_atoms/num_bonds, max_num_bonds) containing the atom or bond indices to select from source.

Returns

A tensor of shape (num_atoms/num_bonds, max_num_bonds, hidden_size) containing the message features corresponding to the atoms/bonds specified in index.

chemprop.nn_utils.initialize_weights(model: torch.nn.modules.module.Module) None[source]

Initializes the weights of a model in place.

Parameters

model – An PyTorch model.

chemprop.nn_utils.param_count(model: torch.nn.modules.module.Module) int[source]

Determines number of trainable parameters.

Parameters

model – An PyTorch model.

Returns

The number of trainable parameters in the model.

chemprop.nn_utils.param_count_all(model: torch.nn.modules.module.Module) int[source]

Determines number of trainable parameters.

Parameters

model – An PyTorch model.

Returns

The number of trainable parameters in the model.