def __init__(self,
num_train_data,
max_updates,
max_epochs,
device="cuda",
fp16_model=False):
config = get_config_with_defaults({
"training": {
"max_updates": max_updates,
"max_epochs": max_epochs,
"evaluation_interval": 10000,
"fp16": True,
},
"run_type": "train",
})
super().__init__(num_train_data, config=config)
if fp16_model:
assert (torch.cuda.is_available()
), "MMFTrainerMock fp16 requires cuda enabled"
model = SimpleModelWithFp16Assert({"in_dim": 1})
model.build()
model = model.cuda()
else:
model = SimpleModel({"in_dim": 1})
model.build()
model.train()
model.to(self.device)
self.model = model
self.optimizer = torch.optim.SGD(self.model.parameters(), lr=1e-3)
def test_eval_loop(self, a, b):
config = get_config_with_defaults(
{"training": {"max_updates": 2, "max_epochs": 2}}
)
trainer = get_mmf_trainer(config=config)
combined_report, meter = trainer.evaluation_loop("val")
self.assertAlmostEqual(combined_report["losses"]["loss"], 493377.5312)
self.assertAlmostEqual(combined_report["logits"].item(), -0.2379742, 6)
def _get_config(self, max_steps, max_epochs):
config = {
"trainer": {
"params": {
"max_steps": max_steps,
"max_epochs": max_epochs
}
}
}
return get_config_with_defaults(config)
def test_loss_computation_parity_with_mmf_trainer(self):
# compute mmf_trainer training losses
def _on_update_end(report, meter, should_log):
self.mmf_losses.append(report["losses"]["loss"].item())
config = get_config_with_defaults(
{"training": {"max_updates": 5, "max_epochs": None}}
)
mmf_trainer = get_mmf_trainer(config=config)
mmf_trainer.on_update_end = _on_update_end
mmf_trainer.evaluation_loop = MagicMock(return_value=(None, None))
mmf_trainer.training_loop()
# compute lightning_trainer training losses
with patch("mmf.trainers.lightning_trainer.get_mmf_env", return_value=""):
config = get_config_with_defaults({"trainer": {"params": {"max_steps": 5}}})
trainer = get_lightning_trainer(config=config)
trainer.callbacks.append(self)
trainer.trainer.fit(trainer.model, trainer.data_module.train_loader)
def _get_config(self, max_steps, max_epochs, gradient_clip_val):
config = {
"trainer": {
"params": {
"max_steps": max_steps,
"max_epochs": max_epochs,
"gradient_clip_val": gradient_clip_val,
}
}
}
return get_config_with_defaults(config)
def _get_mmf_config(self, max_updates, max_epochs, batch_size,
evaluation_interval):
config = {
"training": {
"max_updates": max_updates,
"max_epochs": max_epochs,
"batch_size": batch_size,
"evaluation_interval": evaluation_interval,
}
}
return get_config_with_defaults(config)
def _get_mmf_config(
self, max_updates, max_epochs, max_grad_l2_norm, clip_norm_mode
):
config = {
"training": {
"max_updates": max_updates,
"max_epochs": max_epochs,
"clip_gradients": True,
"max_grad_l2_norm": max_grad_l2_norm,
"clip_norm_mode": clip_norm_mode,
}
}
return get_config_with_defaults(config)
def _get_config(self, accumulate_grad_batches, max_steps, batch_size):
config = {
"trainer": {
"params": {
"accumulate_grad_batches": accumulate_grad_batches,
"max_steps": max_steps,
}
},
"training": {
"batch_size": batch_size
},
}
return get_config_with_defaults(config)
def _get_config(self, max_steps, batch_size, log_every_n_steps,
val_check_interval, tensorboard):
config = {
"trainer": {
"params": {
"max_steps": max_steps,
"log_every_n_steps": log_every_n_steps,
"val_check_interval": val_check_interval,
}
},
"training": {
"batch_size": batch_size,
"tensorboard": tensorboard
},
}
return get_config_with_defaults(config)
def _get_config(
self,
max_updates,
max_epochs,
batch_size=1,
update_frequency=1,
batch_size_per_device=None,
):
config = {
"training": {
"max_updates": max_updates,
"max_epochs": max_epochs,
"update_frequency": update_frequency,
"batch_size": batch_size,
"batch_size_per_device": batch_size_per_device,
}
}
return get_config_with_defaults(config)
def _get_mmf_config(
self,
max_updates,
max_epochs,
batch_size,
log_interval,
evaluation_interval,
tensorboard,
):
config = {
"training": {
"batch_size": batch_size,
"tensorboard": tensorboard,
"max_updates": max_updates,
"max_epochs": max_epochs,
"log_interval": log_interval,
"evaluation_interval": evaluation_interval,
}
}
return get_config_with_defaults(config)
def test_lr_schedule_compared_to_mmf_is_same(self):
config = get_config_with_defaults(
{"training": {"max_updates": 8, "max_epochs": None, "lr_scheduler": True}}
)
mmf_trainer = get_mmf_trainer(config=config)
mmf_trainer.lr_scheduler_callback = LRSchedulerCallback(config, mmf_trainer)
mmf_trainer.callbacks.append(mmf_trainer.lr_scheduler_callback)
mmf_trainer.on_update_end = mmf_trainer.lr_scheduler_callback.on_update_end
mmf_trainer.evaluation_loop = MagicMock(return_value=(None, None))
mmf_trainer.training_loop()
with patch("mmf.trainers.lightning_trainer.get_mmf_env", return_value=""):
config = self._get_config(max_steps=8, lr_scheduler=True)
trainer = get_lightning_trainer(config=config)
trainer.trainer.fit(trainer.model, trainer.data_module.train_loader)
mmf_trainer.model.to(trainer.model.device)
last_model_param1 = list(mmf_trainer.model.parameters())[-1]
last_model_param2 = list(trainer.model.parameters())[-1]
self.assertTrue(torch.allclose(last_model_param1, last_model_param2))
def __init__(self,
config,
num_train_data,
max_updates,
max_epochs,
device="cuda"):
config.training.max_updates = max_updates
config.training.max_epochs = max_epochs
config.training.fp16 = True
config = get_config_with_defaults(config)
super().__init__(num_train_data, config=config)
self.device = torch.device(device)
self.config = config
self.model = SimpleModel({"in_dim": 1})
self.model.build()
self.model = self.model.cuda()
self.optimizer = build_optimizer(self.model, self.config)
self.distributed = True
self.local_rank = 0
self.parallelize_model()
self.load_fp16_scaler()
def _get_config(
self,
max_steps,
batch_size,
val_check_interval,
log_every_n_steps,
limit_val_batches,
):
config = {
"trainer": {
"params": {
"max_steps": max_steps,
"log_every_n_steps": log_every_n_steps,
"val_check_interval": val_check_interval,
"limit_val_batches": limit_val_batches,
}
},
"training": {
"batch_size": batch_size
},
}
return get_config_with_defaults(config)
def _get_config(self, max_steps, lr_scheduler=False):
config = {
"trainer": {"params": {"max_steps": max_steps}},
"training": {"lr_scheduler": lr_scheduler},
}
return get_config_with_defaults(config)
def _get_ckpt_config(self,
is_pl=False,
ckpt_config=None,
max_steps=6,
resume_from_checkpoint=None):
if ckpt_config is None:
ckpt_config = {}
if not is_pl:
return get_config_with_defaults({
"training": {
"max_updates": max_steps,
"max_epochs": None,
"early_stop": {
"enabled": True,
"criteria": "numbers/accuracy",
"minimize": False,
},
"checkpoint_interval": 2,
"evaluation_interval": 2,
},
"model": "simple_model",
"evaluation": {
"metrics": ["accuracy"]
},
"checkpoint": {
"max_to_keep": 1,
"save_git_details": False,
**ckpt_config,
},
"run_type": "train_val",
})
else:
return get_config_with_defaults({
"training": {
"checkpoint_interval": 2,
"early_stop": {
"enabled": True,
"criteria": "numbers/accuracy",
"minimize": False,
},
},
"trainer": {
"params": {
"max_steps": max_steps,
"max_epochs": None,
"checkpoint_callback": True,
"resume_from_checkpoint": resume_from_checkpoint,
"val_check_interval": 2,
}
},
"model": "simple_lightning_model",
"evaluation": {
"metrics": ["accuracy"]
},
"checkpoint": {
"max_to_keep": 1,
"save_git_details": False,
**ckpt_config,
},
"run_type": "train_val",
})
