update

InnerEye/ML/SSL/lightning_containers/ssl_container.py

@@ -157,7 +157,8 @@ def create_model(self) -> LightningModule:
  batch_size=self.data_module.batch_size,
  learning_rate=self.l_rate,
  use_7x7_first_conv_in_resnet=use_7x7_first_conv_in_resnet,
- warmup_epochs=10)
+ warmup_epochs=10,
+ max_epochs=self.num_epochs)
  else:
  raise ValueError(
  f"Unknown value for ssl_training_type, should be {SSLTrainingType.SimCLR.value} or "

InnerEye/ML/SSL/lightning_modules/byol/byol_module.py

@@ -32,6 +32,7 @@ def __init__(self,
  batch_size: int,
  encoder_name: str,
  warmup_epochs: int,
+ max_epochs: int,
  use_7x7_first_conv_in_resnet: bool = True,
  weight_decay: float = 1e-6,
  **kwargs: Any) -> None:

Tests/SSL/byol/test_byol_module.py

@@ -26,7 +26,7 @@ def test_cosine_loss() -> None:
 
 # Test if initial set of parameters are equal between student and teacher.
 def test_module_param_eq() -> None:
- byol = BYOLInnerEye(num_samples=16, learning_rate=1e-3, batch_size=4, encoder_name="resnet50", warmup_epochs=10)
+ byol = BYOLInnerEye(num_samples=16, learning_rate=1e-3, batch_size=4, encoder_name="resnet50", warmup_epochs=10, max_epochs=100)
  pars1 = byol.online_network.parameters()
  pars2 = byol.target_network.parameters()
  for par1, par2 in zip(pars1, pars2):
@@ -39,7 +39,7 @@ def test_encoder_init(encoder_name: str) -> None:
 
 # Test shared step - loss should be bounded between some value and cannot be outside that value.
 def test_shared_forward_step() -> None:
- byol = BYOLInnerEye(num_samples=16, learning_rate=1e-3, batch_size=4, warmup_epochs=10, encoder_name="resnet50")
+ byol = BYOLInnerEye(num_samples=16, learning_rate=1e-3, batch_size=4, warmup_epochs=10, encoder_name="resnet50", max_epochs=100)
  imgs = torch.rand((4, 3, 32, 32))
  lbls = torch.rand((4,))
  batch = ([imgs, imgs], lbls)
@@ -50,7 +50,7 @@ def test_shared_forward_step() -> None:
 
 # Check if output pooling works
 def test_output_spatial_pooling() -> None:
- byol = BYOLInnerEye(num_samples=16, learning_rate=1e-3, batch_size=4, warmup_epochs=10, encoder_name="resnet50")
+ byol = BYOLInnerEye(num_samples=16, learning_rate=1e-3, batch_size=4, warmup_epochs=10, encoder_name="resnet50", max_epochs=100)
  imgs = torch.rand((4, 3, 32, 32))
 
  embeddings = byol(imgs)

Tests/SSL/byol/test_byol_moving_average.py

@@ -43,7 +43,8 @@ def __getitem__(self, item: Any) -> Any:
  learning_rate=1e-3,
  batch_size=4,
  encoder_name="resnet50",
- warmup_epochs=10)
+ warmup_epochs=10,
+ max_epochs=100)
  with mock.patch("InnerEye.ML.SSL.lightning_modules.byol.byol_module.BYOLInnerEye.global_step", 15):
  new_tau = byol_weight_update.update_tau(pl_module=byol_module, trainer=trainer)
  assert new_tau == 1 - 0.01 * (math.cos(math.pi * 15 / total_steps) + 1) / 2