# Copyright 2020 DeepMind Technologies Limited. # # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """BYOL pre-training implementation. Use this experiment to pre-train a self-supervised representation. """ import functools from typing import Any, Generator, Mapping, NamedTuple, Text, Tuple, Union from absl import logging from acme.jax import utils as acme_utils import haiku as hk import jax import jax.numpy as jnp import numpy as np import optax from byol.utils import augmentations from byol.utils import checkpointing from byol.utils import dataset from byol.utils import helpers from byol.utils import networks from byol.utils import optimizers from byol.utils import schedules # Type declarations. LogsDict = Mapping[Text, jnp.ndarray] class _ByolExperimentState(NamedTuple): """Byol's model and optimization parameters and state.""" online_params: hk.Params target_params: hk.Params online_state: hk.State target_state: hk.State opt_state: optimizers.LarsState class ByolExperiment: """Byol's training and evaluation component definition.""" def __init__( self, random_seed: int, num_classes: int, batch_size: int, max_steps: int, enable_double_transpose: bool, base_target_ema: float, network_config: Mapping[Text, Any], optimizer_config: Mapping[Text, Any], lr_schedule_config: Mapping[Text, Any], evaluation_config: Mapping[Text, Any], checkpointing_config: Mapping[Text, Any]): """Constructs the experiment. Args: random_seed: the random seed to use when initializing network weights. num_classes: the number of classes; used for the online evaluation. batch_size: the total batch size; should be a multiple of the number of available accelerators. max_steps: the number of training steps; used for the lr/target network ema schedules. enable_double_transpose: see dataset.py; only has effect on TPU. base_target_ema: the initial value for the ema decay rate of the target network. network_config: the configuration for the network. optimizer_config: the configuration for the optimizer. lr_schedule_config: the configuration for the learning rate schedule. evaluation_config: the evaluation configuration. checkpointing_config: the configuration for checkpointing. """ self._random_seed = random_seed self._enable_double_transpose = enable_double_transpose self._num_classes = num_classes self._lr_schedule_config = lr_schedule_config self._batch_size = batch_size self._max_steps = max_steps self._base_target_ema = base_target_ema self._optimizer_config = optimizer_config self._evaluation_config = evaluation_config # Checkpointed experiment state. self._byol_state = None # Input pipelines. self._train_input = None self._eval_input = None # build the transformed ops forward_fn = functools.partial(self._forward, **network_config) self.forward = hk.without_apply_rng(hk.transform_with_state(forward_fn)) # training can handle multiple devices, thus the pmap self.update_pmap = jax.pmap(self._update_fn, axis_name='i') # evaluation can only handle single device self.eval_batch_jit = jax.jit(self._eval_batch) self._checkpointer = checkpointing.Checkpointer(**checkpointing_config) def _forward( self, inputs: dataset.Batch, projector_hidden_size: int, projector_output_size: int, predictor_hidden_size: int, encoder_class: Text, encoder_config: Mapping[Text, Any], bn_config: Mapping[Text, Any], is_training: bool, ) -> Mapping[Text, jnp.ndarray]: """Forward application of byol's architecture. Args: inputs: A batch of data, i.e. a dictionary, with either two keys, (`images` and `labels`) or three keys (`view1`, `view2`, `labels`). projector_hidden_size: hidden size of the projector MLP. projector_output_size: output size of the projector and predictor MLPs. predictor_hidden_size: hidden size of the predictor MLP. encoder_class: type of the encoder (should match a class in utils/networks). encoder_config: passed to the encoder constructor. bn_config: passed to the hk.BatchNorm constructors. is_training: Training or evaluating the model? When True, inputs must contain keys `view1` and `view2`. When False, inputs must contain key `images`. Returns: All outputs of the model, i.e. a dictionary with projection, prediction and logits keys, for either the two views, or the image. """ encoder = getattr(networks, encoder_class) net = encoder( num_classes=None, # Don't build the final linear layer bn_config=bn_config, **encoder_config) projector = networks.MLP( name='projector', hidden_size=projector_hidden_size, output_size=projector_output_size, bn_config=bn_config) predictor = networks.MLP( name='predictor', hidden_size=predictor_hidden_size, output_size=projector_output_size, bn_config=bn_config) classifier = hk.Linear( output_size=self._num_classes, name='classifier') def apply_once_fn(images: jnp.ndarray, suffix: Text = ''): images = dataset.normalize_images(images) embedding = net(images, is_training=is_training) proj_out = projector(embedding, is_training) pred_out = predictor(proj_out, is_training) # Note the stop_gradient: label information is not leaked into the # main network. classif_out = classifier(jax.lax.stop_gradient(embedding)) outputs = {} outputs['projection' + suffix] = proj_out outputs['prediction' + suffix] = pred_out outputs['logits' + suffix] = classif_out return outputs if is_training: outputs_view1 = apply_once_fn(inputs['view1'], '_view1') outputs_view2 = apply_once_fn(inputs['view2'], '_view2') return {**outputs_view1, **outputs_view2} else: return apply_once_fn(inputs['images'], '') def _optimizer(self, learning_rate: float) -> optax.GradientTransformation: """Build optimizer from config.""" return optimizers.lars( learning_rate, weight_decay_filter=optimizers.exclude_bias_and_norm, lars_adaptation_filter=optimizers.exclude_bias_and_norm, **self._optimizer_config) def loss_fn( self, online_params: hk.Params, target_params: hk.Params, online_state: hk.State, target_state: hk.Params, rng: jnp.ndarray, inputs: dataset.Batch, ) -> Tuple[jnp.ndarray, Tuple[Mapping[Text, hk.State], LogsDict]]: """Compute BYOL's loss function. Args: online_params: parameters of the online network (the loss is later differentiated with respect to the online parameters). target_params: parameters of the target network. online_state: internal state of online network. target_state: internal state of target network. rng: random number generator state. inputs: inputs, containing two batches of crops from the same images, view1 and view2 and labels Returns: BYOL's loss, a mapping containing the online and target networks updated states after processing inputs, and various logs. """ if self._should_transpose_images(): inputs = dataset.transpose_images(inputs) inputs = augmentations.postprocess(inputs, rng) labels = inputs['labels'] online_network_out, online_state = self.forward.apply( params=online_params, state=online_state, inputs=inputs, is_training=True) target_network_out, target_state = self.forward.apply( params=target_params, state=target_state, inputs=inputs, is_training=True) # Representation loss # The stop_gradient is not necessary as we explicitly take the gradient with # respect to online parameters only in `optax.apply_updates`. We leave it to # indicate that gradients are not backpropagated through the target network. repr_loss = helpers.regression_loss( online_network_out['prediction_view1'], jax.lax.stop_gradient(target_network_out['projection_view2'])) repr_loss = repr_loss + helpers.regression_loss( online_network_out['prediction_view2'], jax.lax.stop_gradient(target_network_out['projection_view1'])) repr_loss = jnp.mean(repr_loss) # Classification loss (with gradient flows stopped from flowing into the # ResNet). This is used to provide an evaluation of the representation # quality during training. classif_loss = helpers.softmax_cross_entropy( logits=online_network_out['logits_view1'], labels=jax.nn.one_hot(labels, self._num_classes)) top1_correct = helpers.topk_accuracy( online_network_out['logits_view1'], inputs['labels'], topk=1, ) top5_correct = helpers.topk_accuracy( online_network_out['logits_view1'], inputs['labels'], topk=5, ) top1_acc = jnp.mean(top1_correct) top5_acc = jnp.mean(top5_correct) classif_loss = jnp.mean(classif_loss) loss = repr_loss + classif_loss logs = dict( loss=loss, repr_loss=repr_loss, classif_loss=classif_loss, top1_accuracy=top1_acc, top5_accuracy=top5_acc, ) return loss, (dict(online_state=online_state, target_state=target_state), logs) def _should_transpose_images(self): """Should we transpose images (saves host-to-device time on TPUs).""" return (self._enable_double_transpose and jax.local_devices()[0].platform == 'tpu') def _update_fn( self, byol_state: _ByolExperimentState, global_step: jnp.ndarray, rng: jnp.ndarray, inputs: dataset.Batch, ) -> Tuple[_ByolExperimentState, LogsDict]: """Update online and target parameters. Args: byol_state: current BYOL state. global_step: current training step. rng: current random number generator inputs: inputs, containing two batches of crops from the same images, view1 and view2 and labels Returns: Tuple containing the updated Byol state after processing the inputs, and various logs. """ online_params = byol_state.online_params target_params = byol_state.target_params online_state = byol_state.online_state target_state = byol_state.target_state opt_state = byol_state.opt_state # update online network grad_fn = jax.grad(self.loss_fn, argnums=0, has_aux=True) grads, (net_states, logs) = grad_fn(online_params, target_params, online_state, target_state, rng, inputs) # cross-device grad and logs reductions grads = jax.tree_map(lambda v: jax.lax.pmean(v, axis_name='i'), grads) logs = jax.tree_multimap(lambda x: jax.lax.pmean(x, axis_name='i'), logs) learning_rate = schedules.learning_schedule( global_step, batch_size=self._batch_size, total_steps=self._max_steps, **self._lr_schedule_config) updates, opt_state = self._optimizer(learning_rate).update( grads, opt_state, online_params) online_params = optax.apply_updates(online_params, updates) # update target network tau = schedules.target_ema( global_step, base_ema=self._base_target_ema, max_steps=self._max_steps) target_params = jax.tree_multimap(lambda x, y: x + (1 - tau) * (y - x), target_params, online_params) logs['tau'] = tau logs['learning_rate'] = learning_rate return _ByolExperimentState( online_params=online_params, target_params=target_params, online_state=net_states['online_state'], target_state=net_states['target_state'], opt_state=opt_state), logs def _make_initial_state( self, rng: jnp.ndarray, dummy_input: dataset.Batch, ) -> _ByolExperimentState: """BYOL's _ByolExperimentState initialization. Args: rng: random number generator used to initialize parameters. If working in a multi device setup, this need to be a ShardedArray. dummy_input: a dummy image, used to compute intermediate outputs shapes. Returns: Initial Byol state. """ rng_online, rng_target = jax.random.split(rng) if self._should_transpose_images(): dummy_input = dataset.transpose_images(dummy_input) # Online and target parameters are initialized using different rngs, # in our experiments we did not notice a significant different with using # the same rng for both. online_params, online_state = self.forward.init( rng_online, dummy_input, is_training=True, ) target_params, target_state = self.forward.init( rng_target, dummy_input, is_training=True, ) opt_state = self._optimizer(0).init(online_params) return _ByolExperimentState( online_params=online_params, target_params=target_params, opt_state=opt_state, online_state=online_state, target_state=target_state, ) def step(self, *, global_step: jnp.ndarray, rng: jnp.ndarray) -> Mapping[Text, np.ndarray]: """Performs a single training step.""" if self._train_input is None: self._initialize_train() inputs = next(self._train_input) self._byol_state, scalars = self.update_pmap( self._byol_state, global_step=global_step, rng=rng, inputs=inputs, ) return helpers.get_first(scalars) def save_checkpoint(self, step: int, rng: jnp.ndarray): self._checkpointer.maybe_save_checkpoint( self._byol_state, step=step, rng=rng, is_final=step >= self._max_steps) def load_checkpoint(self) -> Union[Tuple[int, jnp.ndarray], None]: checkpoint_data = self._checkpointer.maybe_load_checkpoint() if checkpoint_data is None: return None self._byol_state, step, rng = checkpoint_data return step, rng def _initialize_train(self): """Initialize train. This includes initializing the input pipeline and Byol's state. """ self._train_input = acme_utils.prefetch(self._build_train_input()) # Check we haven't already restored params if self._byol_state is None: logging.info( 'Initializing parameters rather than restoring from checkpoint.') # initialize Byol and setup optimizer state inputs = next(self._train_input) init_byol = jax.pmap(self._make_initial_state, axis_name='i') # Init uses the same RNG key on all hosts+devices to ensure everyone # computes the same initial state and parameters. init_rng = jax.random.PRNGKey(self._random_seed) init_rng = helpers.bcast_local_devices(init_rng) self._byol_state = init_byol(rng=init_rng, dummy_input=inputs) def _build_train_input(self) -> Generator[dataset.Batch, None, None]: """Loads the (infinitely looping) dataset iterator.""" num_devices = jax.device_count() global_batch_size = self._batch_size per_device_batch_size, ragged = divmod(global_batch_size, num_devices) if ragged: raise ValueError( f'Global batch size {global_batch_size} must be divisible by ' f'num devices {num_devices}') return dataset.load( dataset.Split.TRAIN_AND_VALID, preprocess_mode=dataset.PreprocessMode.PRETRAIN, transpose=self._should_transpose_images(), batch_dims=[jax.local_device_count(), per_device_batch_size]) def _eval_batch( self, params: hk.Params, state: hk.State, batch: dataset.Batch, ) -> Mapping[Text, jnp.ndarray]: """Evaluates a batch. Args: params: Parameters of the model to evaluate. Typically Byol's online parameters. state: State of the model to evaluate. Typically Byol's online state. batch: Batch of data to evaluate (must contain keys images and labels). Returns: Unreduced evaluation loss and top1 accuracy on the batch. """ if self._should_transpose_images(): batch = dataset.transpose_images(batch) outputs, _ = self.forward.apply(params, state, batch, is_training=False) logits = outputs['logits'] labels = hk.one_hot(batch['labels'], self._num_classes) loss = helpers.softmax_cross_entropy(logits, labels, reduction=None) top1_correct = helpers.topk_accuracy(logits, batch['labels'], topk=1) top5_correct = helpers.topk_accuracy(logits, batch['labels'], topk=5) # NOTE: Returned values will be summed and finally divided by num_samples. return { 'eval_loss': loss, 'top1_accuracy': top1_correct, 'top5_accuracy': top5_correct, } def _eval_epoch(self, subset: Text, batch_size: int): """Evaluates an epoch.""" num_samples = 0. summed_scalars = None params = helpers.get_first(self._byol_state.online_params) state = helpers.get_first(self._byol_state.online_state) split = dataset.Split.from_string(subset) dataset_iterator = dataset.load( split, preprocess_mode=dataset.PreprocessMode.EVAL, transpose=self._should_transpose_images(), batch_dims=[batch_size]) for inputs in dataset_iterator: num_samples += inputs['labels'].shape[0] scalars = self.eval_batch_jit(params, state, inputs) # Accumulate the sum of scalars for each step. scalars = jax.tree_map(lambda x: jnp.sum(x, axis=0), scalars) if summed_scalars is None: summed_scalars = scalars else: summed_scalars = jax.tree_multimap(jnp.add, summed_scalars, scalars) mean_scalars = jax.tree_map(lambda x: x / num_samples, summed_scalars) return mean_scalars def evaluate(self, global_step, **unused_args): """Thin wrapper around _eval_epoch.""" global_step = np.array(helpers.get_first(global_step)) scalars = jax.device_get(self._eval_epoch(**self._evaluation_config)) logging.info('[Step %d] Eval scalars: %s', global_step, scalars) return scalars