def load_model_and_optimizer(self):
attributes = self.config.model_config[self.config.model]
# Easy way to point to config for other model
if isinstance(attributes, str):
attributes = self.config.model_config[attributes]
with omegaconf.open_dict(attributes):
attributes.model = self.config.model
self.model = build_model(attributes)
if "cuda" in str(self.device):
device_info = "CUDA Device {} is: {}".format(
self.config.distributed.rank,
torch.cuda.get_device_name(self.local_rank),
)
registry.register("global_device", self.config.distributed.rank)
self.writer.write(device_info, log_all=True)
self.model = self.model.to(self.device)
self.optimizer = build_optimizer(self.model, self.config)
registry.register("data_parallel", False)
registry.register("distributed", False)
self.load_extras()
self.parallelize_model()
def get_mmf_trainer(
model_size=1,
num_data_size=100,
max_updates=5,
max_epochs=None,
on_update_end_fn=None,
fp16=False,
scheduler_config=None,
grad_clipping_config=None,
):
torch.random.manual_seed(2)
model = SimpleModel(model_size)
model.train()
trainer_config = get_trainer_config()
optimizer = build_optimizer(model, trainer_config)
trainer = TrainerTrainingLoopMock(
num_data_size,
max_updates,
max_epochs,
config=trainer_config,
optimizer=optimizer,
on_update_end_fn=on_update_end_fn,
fp16=fp16,
scheduler_config=scheduler_config,
grad_clipping_config=grad_clipping_config,
)
model.to(trainer.device)
trainer.model = model
return trainer
def test_build_optimizer_custom_model(self):
model = MMBT.from_params()
model.build()
self.config.model = model.config.model
self.config.model_config = model.config
optimizer = build_optimizer(model, self.config)
self.assertTrue(isinstance(optimizer, torch.optim.Optimizer))
self.assertEqual(len(optimizer.param_groups), 2)
def __init__(
self,
config,
num_train_data,
max_updates,
max_epochs,
device="cuda",
fp16_model=False,
):
super().__init__(num_train_data, max_updates, max_epochs, fp16=True)
self.device = torch.device(device)
self.config = config
self.model = SimpleModel(1)
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 __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_mmf_trainer(
model_size=1,
num_data_size=100,
max_updates=5,
max_epochs=None,
on_update_end_fn=None,
fp16=False,
scheduler_config=None,
grad_clipping_config=None,
evaluation_interval=4,
log_interval=1,
batch_size=1,
tensorboard=False,
):
torch.random.manual_seed(2)
model = SimpleModel({"in_dim": model_size})
model.build()
model.train()
trainer_config = get_trainer_config()
trainer_config.training.evaluation_interval = evaluation_interval
trainer_config.training.log_interval = log_interval
optimizer = build_optimizer(model, trainer_config)
trainer = TrainerTrainingLoopMock(
num_data_size,
max_updates,
max_epochs,
config=trainer_config,
optimizer=optimizer,
fp16=fp16,
on_update_end_fn=on_update_end_fn,
scheduler_config=scheduler_config,
grad_clipping_config=grad_clipping_config,
batch_size=batch_size,
tensorboard=tensorboard,
)
trainer.load_datasets()
model.to(trainer.device)
trainer.model = model
return trainer
def load_optimizer(self):
logger.info("Loading optimizer")
self.optimizer = build_optimizer(self.model, self.config)
def test_build_optimizer_simple_model(self):
model = SimpleModel({"in_dim": 1})
model.build()
optimizer = build_optimizer(model, self.config)
self.assertTrue(isinstance(optimizer, torch.optim.Optimizer))
self.assertEqual(len(optimizer.param_groups), 1)
def load_optimizer(self):
self.writer.write("Loading optimizer", "info")
self.optimizer = build_optimizer(self.model, self.config)
def add_optimizer(trainer, config):
optimizer = build_optimizer(trainer.model, config)
trainer.optimizer = optimizer
