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)
