BYOL improvements

configs/config/benchmark/linear_image_classification/imagenet1k/byol_transfer_in1k_linear.yaml

@@ -22,6 +22,7 @@ config:
  TRANSFORMS:
  - name: RandomResizedCrop
  size: 224
+ interpolation: 3
  - name: RandomHorizontalFlip
  - name: ToTensor
  - name: Normalize
@@ -38,6 +39,7 @@ config:
  TRANSFORMS:
  - name: Resize
  size: 256
+ interpolation: 3
  - name: CenterCrop
  size: 224
  - name: ToTensor
@@ -82,7 +84,7 @@ config:
  PARAMS_FILE: "specify the model weights"
  STATE_DICT_KEY_NAME: classy_state_dict
  SYNC_BN_CONFIG:
- CONVERT_BN_TO_SYNC_BN: True
+ CONVERT_BN_TO_SYNC_BN: False
  SYNC_BN_TYPE: apex
  GROUP_SIZE: 8
  LOSS:
@@ -93,22 +95,29 @@ config:
  name: sgd
  momentum: 0.9
  num_epochs: 80
+ weight_decay: 0
  nesterov: True
  regularize_bn: False
  regularize_bias: True
  param_schedulers:
  lr:
  auto_lr_scaling:
- auto_scale: true
- base_value: 0.4
+ # if set to True, learning rate will be scaled.
+ auto_scale: True
+ # base learning rate value that will be scaled.
+ base_value: 0.2
+ # batch size for which the base learning rate is specified. The current batch size
+ # is used to determine how to scale the base learning rate value.
+ # scaled_lr = ((batchsize_per_gpu * world_size) * base_value ) / base_lr_batch_size
  base_lr_batch_size: 256
- name: multistep
- values: [0.4, 0.3, 0.2, 0.1, 0.05]
- milestones: [16, 32, 48, 64]
- update_interval: epoch
+ # scaling_type can be set to "sqrt" to reduce the impact of scaling on the base value
+ scaling_type: "linear"
+ name: constant
+ update_interval: "epoch"
+ value: 0.2
  DISTRIBUTED:
  BACKEND: nccl
- NUM_NODES: 8
+ NUM_NODES: 4
  NUM_PROC_PER_NODE: 8
  INIT_METHOD: tcp
  RUN_ID: auto

configs/config/pretrain/byol/byol_8node_resnet.yaml

@@ -67,7 +67,7 @@ config:
  RESNETS:
  DEPTH: 50
  ZERO_INIT_RESIDUAL: True
- HEAD: 
+ HEAD:
  PARAMS: [
  ["mlp", {"dims": [2048, 4096, 256], "use_relu": True, "use_bn": True}],
  ["mlp", {"dims": [256, 4096, 256], "use_relu": True, "use_bn": True}]
@@ -82,15 +82,16 @@ config:
  byol_loss:
  embedding_dim: 256
  momentum: 0.99
- OPTIMIZER: # from official BYOL implementation, deepmind-research/byol/configs/byol.py
+ OPTIMIZER:
  name: lars
- trust_coefficient: 0.001
+ eta: 0.001
  weight_decay: 1.0e-6
  momentum: 0.9
  nesterov: False
  num_epochs: 300
  regularize_bn: False
- regularize_bias: True
+ regularize_bias: False
+ exclude_bias_and_norm: True
  param_schedulers:
  lr:
  auto_lr_scaling:

vissl/data/ssl_transforms/img_pil_color_distortion.py

@@ -21,8 +21,16 @@ class ImgPilColorDistortion(ClassyTransform):
  randomly convert the image to grayscale.
  """
 
- def __init__(self, strength, brightness=0.8, contrast=0.8, saturation=0.8,
- hue=0.2, color_jitter_probability=0.8, grayscale_probability=0.2):
+ def __init__(
+ self,
+ strength,
+ brightness=0.8,
+ contrast=0.8,
+ saturation=0.8,
+ hue=0.2,
+ color_jitter_probability=0.8,
+ grayscale_probability=0.2,
+ ):
  """
  Args:
  strength (float): A number used to quantify the strength of the
@@ -41,22 +49,23 @@ def __init__(self, strength, brightness=0.8, contrast=0.8, saturation=0.8,
  grayscale_probability (float): A floating point number used to
  quantify to apply randomly convert image to grayscale with
  the assigned probability. Default value is 0.2.
- This function follows the Pytorch documentation: https://pytorch.org/vision/stable/transforms.html
  """
  self.strength = strength
  self.brightness = brightness
  self.contrast = contrast
  self.saturation = saturation
  self.hue = hue
- self.color_jitter_probability=color_jitter_probability
- self.grayscale_probability=grayscale_probability
+ self.color_jitter_probability = color_jitter_probability
+ self.grayscale_probability = grayscale_probability
  self.color_jitter = pth_transforms.ColorJitter(
  self.brightness * self.strength,
  self.contrast * self.strength,
  self.saturation * self.strength,
  self.hue * self.strength,
  )
- self.rnd_color_jitter = pth_transforms.RandomApply([self.color_jitter], p=self.color_jitter_probability)
+ self.rnd_color_jitter = pth_transforms.RandomApply(
+ [self.color_jitter], p=self.color_jitter_probability
+ )
  self.rnd_gray = pth_transforms.RandomGrayscale(p=self.grayscale_probability)
  self.transforms = pth_transforms.Compose([self.rnd_color_jitter, self.rnd_gray])
 

vissl/hooks/__init__.py

@@ -8,6 +8,7 @@
 
 from classy_vision.hooks.classy_hook import ClassyHook
 from vissl.config import AttrDict
+from vissl.hooks.byol_hooks import BYOLHook # noqa
 from vissl.hooks.deepclusterv2_hooks import ClusterMemoryHook, InitMemoryHook # noqa
 from vissl.hooks.dino_hooks import DINOHook
 from vissl.hooks.grad_clip_hooks import GradClipHook # noqa
@@ -21,15 +22,12 @@
 )
 from vissl.hooks.moco_hooks import MoCoHook # noqa
 from vissl.hooks.profiling_hook import ProfilingHook
-from vissl.hooks.byol_hooks import BYOLHook # noqa
-
 from vissl.hooks.state_update_hooks import ( # noqa
  CheckNanLossHook,
  FreezeParametersHook,
  SetDataSamplerEpochHook,
  SSLModelComplexityHook,
 )
-from vissl.hooks.byol_hooks import BYOLHook # noqa
 from vissl.hooks.swav_hooks import NormalizePrototypesHook # noqa
 from vissl.hooks.swav_hooks import SwAVUpdateQueueScoresHook # noqa
 from vissl.hooks.swav_momentum_hooks import (
@@ -149,14 +147,6 @@ def default_hook_generator(cfg: AttrDict) -> List[ClassyHook]:
  )
  ]
  )
- if cfg.LOSS.name == "byol_loss":
- hooks.extend(
- [
- BYOLHook(
- cfg.LOSS["byol_loss"]["momentum"],
- )
- ]
- )
  if cfg.HOOKS.MODEL_COMPLEXITY.COMPUTE_COMPLEXITY:
  hooks.extend([SSLModelComplexityHook()])
  if cfg.HOOKS.LOG_GPU_STATS:

vissl/hooks/byol_hooks.py

@@ -1,21 +1,22 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-import math
 import logging
+import math
 
 import torch
 from classy_vision import tasks
 from classy_vision.hooks.classy_hook import ClassyHook
 from vissl.models import build_model
 from vissl.utils.env import get_machine_local_and_dist_rank
 
+
 class BYOLHook(ClassyHook):
  """
- BYOL - Bootstrap your own latent: (https://arxiv.org/abs/2006.07733)
- is based on Contrastive learning. This hook
- creates a target network with the same architecture
- as the main online network, but without the projection head.
- The online network does not participate in backpropogation,
- but instead is an exponential moving average of the online network.
+  BYOL - Bootstrap your own latent: (https://arxiv.org/abs/2006.07733)
+  is based on Contrastive learning. This hook
+  creates a target network with the same architecture
+  as the main online network, but without the projection head.
+  The online network does not participate in backpropogation,
+  but instead is an exponential moving average of the online network.
  """
 
  on_start = ClassyHook._noop
@@ -28,7 +29,7 @@ class BYOLHook(ClassyHook):
  on_update = ClassyHook._noop
 
  @staticmethod
- def cosine_decay(training_iter, max_iters, initial_value)  -> float:
+ def cosine_decay(training_iter, max_iters, initial_value) -> float:
  """
  For a given starting value, this function anneals the learning
  rate.
@@ -42,8 +43,8 @@ def target_ema(training_iter, base_ema, max_iters) -> float:
  """
  Updates Exponential Moving average of the Target Network.
  """
- decay = BYOLHook.cosine_decay(training_iter, max_iters, 1.)
- return 1. - (1. - base_ema) * decay
+ decay = BYOLHook.cosine_decay(training_iter, max_iters, 1.0)
+ return 1.0 - (1.0 - base_ema) * decay
 
  def _build_byol_target_network(self, task: tasks.ClassyTask) -> None:
  """
@@ -53,27 +54,29 @@ def _build_byol_target_network(self, task: tasks.ClassyTask) -> None:
  """
  # Create the encoder, which will slowly track the model
  logging.info(
- "BYOL: Building BYOL target network - rank %s %s", *get_machine_local_and_dist_rank()
+ "BYOL: Building BYOL target network - rank %s %s",
+ *get_machine_local_and_dist_rank(),
  )
 
- # Target model has the same architecture, but without the projector head.
- target_model_config = task.config['MODEL']
- target_model_config['HEAD']['PARAMS'] = target_model_config['HEAD']['PARAMS'][0:1]
+ # Target model has the same architecture, *without* the projector head.
+ target_model_config = task.config["MODEL"]
+ target_model_config["HEAD"]["PARAMS"] = target_model_config["HEAD"]["PARAMS"][
+ 0:1
+ ]
  task.loss.target_network = build_model(
  target_model_config, task.config["OPTIMIZER"]
  )
 
- # TESTED: Target Network and Online network are properly created.
- # TODO: Check SyncBatchNorm settings (low prior)
-
  task.loss.target_network.to(task.device)
 
  # Restore an hypothetical checkpoint, else copy the model parameters from the
  # online network.
  if task.loss.checkpoint is not None:
  task.loss.load_state_dict(task.loss.checkpoint)
  else:
- logging.info("BYOL: Copying and freezing model parameters from online to target network")
+ logging.info(
+ "BYOL: Copying and freezing model parameters from online to target network"
+ )
  for param_q, param_k in zip(
  task.base_model.parameters(), task.loss.target_network.parameters()
  ):
@@ -92,7 +95,9 @@ def _update_momentum_coefficient(self, task: tasks.ClassyTask) -> None:
  self.total_iters = task.max_iteration
  logging.info(f"{self.total_iters} total iters")
  training_iteration = task.iteration
- self.momentum = self.target_ema(training_iteration, self.base_momentum, self.total_iters)
+ self.momentum = self.target_ema(
+ training_iteration, self.base_momentum, self.total_iters
+ )
 
  @torch.no_grad()
  def _update_target_network(self, task: tasks.ClassyTask) -> None:
@@ -106,10 +111,10 @@ def _update_target_network(self, task: tasks.ClassyTask) -> None:
  task.base_model.parameters(), task.loss.target_network.parameters()
  ):
  target_params.data = (
- target_params.data * self.momentum + online_params.data * (1. - self.momentum)
+ target_params.data * self.momentum
+ + online_params.data * (1.0 - self.momentum)
  )
 
-
  @torch.no_grad()
  def on_forward(self, task: tasks.ClassyTask) -> None:
  """
@@ -127,9 +132,8 @@ def on_forward(self, task: tasks.ClassyTask) -> None:
  else:
  self._update_target_network(task)
 
-
  # Compute target network embeddings
- batch = task.last_batch.sample['input']
+ batch = task.last_batch.sample["input"]
  target_embs = task.loss.target_network(batch)[0]
 
  # Save target embeddings to use them in the loss

vissl/losses/byol_loss.py

@@ -7,9 +7,9 @@
 import torch.nn.functional as F
 from classy_vision.losses import ClassyLoss, register_loss
 
-_BYOLLossConfig = namedtuple(
- "_BYOLLossConfig", ["embedding_dim", "momentum"]
-)
+
+_BYOLLossConfig = namedtuple("_BYOLLossConfig", ["embedding_dim", "momentum"])
+
 
 def regression_loss(x, y):
  """
@@ -19,17 +19,16 @@ def regression_loss(x, y):
  Cosine similarity. This implementation uses Cosine similarity.
  """
  normed_x, normed_y = F.normalize(x, dim=1), F.normalize(y, dim=1)
- return torch.sum((normed_x - normed_y).pow(2), dim=1)
+ # Euclidean Distance squared.
+ return 2 - 2 * (normed_x * normed_y).sum(dim=1)
 
 
 class BYOLLossConfig(_BYOLLossConfig):
- """ Settings for the BYOL loss"""
+ """Settings for the BYOL loss"""
 
  @staticmethod
  def defaults() -> "BYOLLossConfig":
- return BYOLLossConfig(
- embedding_dim=256, momentum=0.999
- )
+ return BYOLLossConfig(embedding_dim=256, momentum=0.999)
 
 
 @register_loss("byol_loss")
@@ -68,7 +67,9 @@ def from_config(cls, config: BYOLLossConfig) -> "BYOLLoss":
  """
  return cls(config)
 
- def forward(self, online_network_prediction: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ def forward(
+ self, online_network_prediction: torch.Tensor, *args, **kwargs
+ ) -> torch.Tensor:
  """
  In this function, the Online Network receives the tensor as input after projection
  and they make predictions on the output of the target network’s projection,
@@ -79,7 +80,6 @@ def forward(self, online_network_prediction: torch.Tensor, *args, **kwargs) -> t
  compute the cross entropy loss for this batch.
 
  Args:
- query: output of the encoder given the current batch
  online_network_prediction: online model output. this is a prediction of the
  target network output.
 
@@ -91,8 +91,6 @@ def forward(self, online_network_prediction: torch.Tensor, *args, **kwargs) -> t
  online_view1, online_view2 = torch.chunk(online_network_prediction, 2, 0)
  target_view1, target_view2 = torch.chunk(self.target_embs.detach(), 2, 0)
 
- # TESTED: Views are received correctly.
-
  # Compute losses
  loss1 = regression_loss(online_view1, target_view2)
  loss2 = regression_loss(online_view2, target_view1)
@@ -111,7 +109,6 @@ def load_state_dict(self, state_dict, *args, **kwargs) -> None:
  Args:
  state_dict (serialized via torch.save)
  """
-
  # If the encoder has been allocated, use the normal pytorch restoration
  if self.target_network is None:
  self.checkpoint = state_dict