Merge pull request #19 from RMeli/jit

Compile models and losses with TorchScript for training and inference

gnina/inference.py

@@ -153,6 +153,9 @@ def inference(args):
  # Create model
  model = models.models_dict[(args.model, affinity)](test_loader.dims).to(device)
 
+ # Compile model with TorchScript
+ model = torch.jit.script(model)
+
  # Load checkpoint
  checkpoint = torch.load(args.checkpoint, map_location=device)
  Checkpoint.load_objects(to_load={"model": model}, checkpoint=checkpoint)

gnina/losses.py

@@ -90,7 +90,7 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor:
 
  if self.reduction == "mean":
  reduced_loss = torch.mean(loss)
- elif self.reduction == "sum":
+ else: # Assertion in init ensures that reduction is "sum"
  reduced_loss = torch.sum(loss)
 
  return self.scale * reduced_loss

gnina/training.py

@@ -589,8 +589,8 @@ def training(args):
  model = models_dict[(args.model, affinity)](train_loader.dims).to(device)
  model.apply(weights_and_biases_init)
 
- # TODO: Compile model into TorchScript
- # Requires model refactoring to avoid branching based on affinity
+ # Compile model into TorchScript
+ model = torch.jit.script(model)
 
  optimizer = optim.SGD(
  model.parameters(),
@@ -600,13 +600,15 @@ def training(args):
  )
 
  # Define loss functions
- pose_loss = nn.NLLLoss()
+ pose_loss = torch.jit.script(nn.NLLLoss())
  affinity_loss = (
- AffinityLoss(
- delta=args.delta_affinity_loss,
- penalty=args.penalty_affinity_loss,
- pseudo_huber=args.pseudo_huber_affinity_loss,
- scale=args.scale_affinity_loss,
+ torch.jit.script(
+ AffinityLoss(
+ delta=args.delta_affinity_loss,
+ penalty=args.penalty_affinity_loss,
+ pseudo_huber=args.pseudo_huber_affinity_loss,
+ scale=args.scale_affinity_loss,
+ )
  )
  if affinity
  else None