def _compute_param_groups(finetune_config: AttrDict): """ Take a configuration and compute the parameter groups for this configuration """ optimizer_schedulers = build_optimizer_schedulers( finetune_config["OPTIMIZER"]) base_model = build_model(finetune_config["MODEL"], finetune_config["OPTIMIZER"]) return get_optimizer_param_groups( model=base_model, model_config=finetune_config["MODEL"], optimizer_config=finetune_config["OPTIMIZER"], optimizer_schedulers=optimizer_schedulers, )
def _build_task(self, num_epochs, skip_param_schedulers=False): config = self._get_config(skip_param_schedulers) config["optimizer"]["num_epochs"] = num_epochs task = (ClassificationTask().set_num_epochs(num_epochs).set_loss( build_loss(config["loss"])).set_model(build_model( config["model"])).set_optimizer( build_optimizer( config["optimizer"])).set_optimizer_schedulers( build_optimizer_schedulers(config["optimizer"]))) for phase_type in ["train", "test"]: dataset = build_dataset(config["dataset"][phase_type]) task.set_dataset(dataset, phase_type) self.assertTrue(task is not None) return task
def test_training(self): """Checks we can train a small MLP model.""" config = get_test_mlp_task_config() task = (ClassificationTask().set_num_epochs(10).set_loss( build_loss(config["loss"])).set_model(build_model( config["model"])).set_optimizer( build_optimizer( config["optimizer"])).set_optimizer_schedulers( build_optimizer_schedulers( config["optimizer"])).set_meters([ AccuracyMeter(topk=[1]) ]).set_hooks([LossLrMeterLoggingHook()])) for split in ["train", "test"]: dataset = build_dataset(config["dataset"][split]) task.set_dataset(dataset, split) self.assertTrue(task is not None) trainer = LocalTrainer() trainer.train(task) accuracy = task.meters[0].value["top_1"] self.assertAlmostEqual(accuracy, 1.0)
def _build_optimizer_schedulers(self): """ Build the param schedulers to be used in training. """ return build_optimizer_schedulers(self.config["OPTIMIZER"])
def from_config(cls, config: Dict[str, Any]) -> "ClassificationTask": """Instantiates a ClassificationTask from a configuration. Args: config: A configuration for a ClassificationTask. See :func:`__init__` for parameters expected in the config. Returns: A ClassificationTask instance. """ test_only = config.get("test_only", False) if not test_only: # TODO Make distinction between epochs and phases in optimizer clear train_phases_per_epoch = config["dataset"]["train"].get( "phases_per_epoch", 1) optimizer_config = config["optimizer"] optimizer_config["num_epochs"] = (config["num_epochs"] * train_phases_per_epoch) optimizer = build_optimizer(optimizer_config) param_schedulers = build_optimizer_schedulers(optimizer_config) datasets = {} phase_types = ["train", "test"] for phase_type in phase_types: if phase_type in config["dataset"]: datasets[phase_type] = build_dataset( config["dataset"][phase_type]) loss = build_loss(config["loss"]) amp_args = config.get("amp_args") meters = build_meters(config.get("meters", {})) model = build_model(config["model"]) mixup_transform = None if config.get("mixup") is not None: assert "alpha" in config[ "mixup"], "key alpha is missing in mixup dict" mixup_transform = MixupTransform( config["mixup"]["alpha"], config["mixup"].get("num_classes")) # hooks config is optional hooks_config = config.get("hooks") hooks = [] if hooks_config is not None: hooks = build_hooks(hooks_config) distributed_config = config.get("distributed", {}) distributed_options = { "broadcast_buffers_mode": BroadcastBuffersMode[distributed_config.get( "broadcast_buffers", "before_eval").upper()], "batch_norm_sync_mode": BatchNormSyncMode[distributed_config.get("batch_norm_sync_mode", "disabled").upper()], "batch_norm_sync_group_size": distributed_config.get("batch_norm_sync_group_size", 0), "find_unused_parameters": distributed_config.get("find_unused_parameters", True), } task = ( cls().set_num_epochs(config["num_epochs"]).set_test_phase_period( config.get( "test_phase_period", 1)).set_loss(loss).set_test_only(test_only).set_model( model).set_meters(meters).set_amp_args(amp_args). set_mixup_transform(mixup_transform).set_distributed_options( **distributed_options).set_hooks(hooks).set_bn_weight_decay( config.get("bn_weight_decay", False))) if not test_only: task.set_optimizer(optimizer) task.set_optimizer_schedulers(param_schedulers) use_gpu = config.get("use_gpu") if use_gpu is not None: task.set_use_gpu(use_gpu) for phase_type in datasets: task.set_dataset(datasets[phase_type], phase_type) # NOTE: this is a private member and only meant to be used for # logging/debugging purposes. See __repr__ implementation task._config = config return task
def test_lr_schedule(self): config = self._get_config() opt = build_optimizer(config) param_schedulers = build_optimizer_schedulers(config) opt.set_param_groups({"params": self._parameters(), **param_schedulers}) # Test initial learning rate for group in opt.optimizer.param_groups: self.assertEqual(group["lr"], 0.1) def _test_lr_schedule(optimizer, num_epochs, epochs, targets): for i in range(len(epochs)): epoch = epochs[i] target = targets[i] param_groups = optimizer.optimizer.param_groups.copy() optimizer.on_epoch(epoch / num_epochs) for idx, group in enumerate(optimizer.optimizer.param_groups): self.assertEqual(group["lr"], target) # Make sure all but LR is same param_groups[idx]["lr"] = target self.assertEqual(param_groups[idx], group) # Test constant learning schedule num_epochs = 90 epochs = [0, 0.025, 0.05, 0.1, 0.5, 1, 15, 29, 30, 31, 59, 60, 61, 88, 89] targets = [0.1] * 15 _test_lr_schedule(opt, num_epochs, epochs, targets) # Test step learning schedule config["param_schedulers"] = { "lr": {"name": "step", "values": [0.1, 0.01, 0.001]} } opt = build_optimizer(config) param_schedulers = build_optimizer_schedulers(config) opt.set_param_groups({"params": self._parameters(), **param_schedulers}) targets = [0.1] * 8 + [0.01] * 3 + [0.001] * 4 _test_lr_schedule(opt, num_epochs, epochs, targets) # Test step learning schedule with warmup init_lr = 0.01 warmup_epochs = 0.1 config["param_schedulers"] = { "lr": { "name": "composite", "schedulers": [ {"name": "linear", "start_value": init_lr, "end_value": 0.1}, {"name": "step", "values": [0.1, 0.01, 0.001]}, ], "update_interval": "epoch", "interval_scaling": ["rescaled", "fixed"], "lengths": [warmup_epochs / num_epochs, 1 - warmup_epochs / num_epochs], } } opt = build_optimizer(config) param_schedulers = build_optimizer_schedulers(config) opt.set_param_groups({"params": self._parameters(), **param_schedulers}) targets = [0.01, 0.0325, 0.055] + [0.1] * 5 + [0.01] * 3 + [0.001] * 4 _test_lr_schedule(opt, num_epochs, epochs, targets)