class TrainerTrainingLoopMock(TrainerTrainingLoopMixin, TrainerProfilingMixin):
def __init__(self, num_train_data, max_updates, max_epochs):
self.training_config = OmegaConf.create({
"detect_anomaly": False,
"evaluation_interval": 10000
})
if max_updates is not None:
self.training_config["max_updates"] = max_updates
if max_epochs is not None:
self.training_config["max_epochs"] = max_epochs
self.model = SimpleModel(1)
if torch.cuda.is_available():
self.model = self.model.cuda()
self.dataset_loader = MagicMock()
self.dataset_loader.seed_sampler = MagicMock(return_value=None)
self.dataset_loader.prepare_batch = lambda x: SampleList(x)
self.optimizer = MagicMock()
self.optimizer.step = MagicMock(return_value=None)
self.optimizer.zero_grad = MagicMock(return_value=None)
dataset = NumbersDataset(num_train_data)
self.train_loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=1,
shuffle=False,
num_workers=1,
drop_last=False,
)
self.on_batch_start = MagicMock(return_value=None)
self.logistics_callback = MagicMock(return_value=None)
self.logistics_callback.log_interval = MagicMock(return_value=None)
self.on_batch_end = MagicMock(return_value=None)
self.meter = MagicMock(return_value=None)
self.after_training_loop = MagicMock(return_value=None)
class MMFTrainerMock(TrainerTrainingLoopMock, MMFTrainer):
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()
class MMFTrainerMock(TrainerTrainingLoopMock, MMFTrainer):
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()
class TrainerTrainingLoopMock(MMFTrainer):
def __init__(
self,
num_train_data,
max_updates,
max_epochs,
config=None,
optimizer=None,
update_frequency=1,
batch_size=1,
batch_size_per_device=None,
fp16=False,
on_update_end_fn=None,
scheduler_config=None,
grad_clipping_config=None,
):
if config is None:
self.config = OmegaConf.create(
{
"training": {
"detect_anomaly": False,
"evaluation_interval": 10000,
"update_frequency": update_frequency,
"fp16": fp16,
"batch_size": batch_size,
"batch_size_per_device": batch_size_per_device,
}
}
)
self.training_config = self.config.training
else:
self.training_config = config.training
self.config = config
# Load batch size with custom config and cleanup
original_config = registry.get("config")
registry.register("config", self.config)
batch_size = get_batch_size()
registry.register("config", original_config)
if max_updates is not None:
self.training_config["max_updates"] = max_updates
if max_epochs is not None:
self.training_config["max_epochs"] = max_epochs
self.model = SimpleModel({"in_dim": 1})
self.model.build()
if torch.cuda.is_available():
self.model = self.model.cuda()
self.device = "cuda"
else:
self.device = "cpu"
self.distributed = False
self.dataset_loader = MagicMock()
self.dataset_loader.seed_sampler = MagicMock(return_value=None)
self.dataset_loader.prepare_batch = lambda x: SampleList(x)
if optimizer is None:
self.optimizer = MagicMock()
self.optimizer.step = MagicMock(return_value=None)
self.optimizer.zero_grad = MagicMock(return_value=None)
else:
self.optimizer = optimizer
if scheduler_config:
config.training.lr_scheduler = True
config.scheduler = scheduler_config
self.lr_scheduler_callback = LRSchedulerCallback(config, self)
self.callbacks.append(self.lr_scheduler_callback)
on_update_end_fn = (
on_update_end_fn
if on_update_end_fn
else self.lr_scheduler_callback.on_update_end
)
if grad_clipping_config:
self.training_config.clip_gradients = True
self.training_config.max_grad_l2_norm = grad_clipping_config[
"max_grad_l2_norm"
]
self.training_config.clip_norm_mode = grad_clipping_config["clip_norm_mode"]
dataset = NumbersDataset(num_train_data)
self.train_loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=batch_size,
shuffle=False,
num_workers=1,
drop_last=False,
)
self.train_loader.current_dataset = dataset
self.on_batch_start = MagicMock(return_value=None)
self.on_update_start = MagicMock(return_value=None)
self.logistics_callback = MagicMock(return_value=None)
self.logistics_callback.log_interval = MagicMock(return_value=None)
self.on_batch_end = MagicMock(return_value=None)
self.on_update_end = (
on_update_end_fn if on_update_end_fn else MagicMock(return_value=None)
)
self.meter = Meter()
self.after_training_loop = MagicMock(return_value=None)
self.on_validation_start = MagicMock(return_value=None)
self.evaluation_loop = MagicMock(return_value=(None, None))
self.scaler = torch.cuda.amp.GradScaler(enabled=False)
self.val_loader = MagicMock(return_value=None)
self.early_stop_callback = MagicMock(return_value=None)
self.on_validation_end = MagicMock(return_value=None)
self.metrics = MagicMock(return_value=None)
class TrainerTrainingLoopMock(MMFTrainer):
def __init__(
self,
num_train_data,
max_updates,
max_epochs,
config=None,
optimizer=None,
update_frequency=1,
batch_size=1,
batch_size_per_device=None,
fp16=False,
on_update_end_fn=None,
scheduler_config=None,
grad_clipping_config=None,
tensorboard=False,
):
if config is None:
self.config = OmegaConf.create({
"training": {
"detect_anomaly": False,
"evaluation_interval": 10000,
"update_frequency": update_frequency,
"fp16": fp16,
"batch_size": batch_size,
"batch_size_per_device": batch_size_per_device,
"tensorboard": tensorboard,
}
})
self.training_config = self.config.training
else:
config.training.batch_size = batch_size
config.training.fp16 = fp16
config.training.update_frequency = update_frequency
config.training.tensorboard = tensorboard
self.training_config = config.training
self.config = config
registry.register("config", self.config)
if max_updates is not None:
self.training_config["max_updates"] = max_updates
if max_epochs is not None:
self.training_config["max_epochs"] = max_epochs
self.model = SimpleModel({"in_dim": 1})
self.model.build()
if torch.cuda.is_available():
self.model = self.model.cuda()
self.device = "cuda"
else:
self.device = "cpu"
self.distributed = False
if optimizer is None:
self.optimizer = MagicMock()
self.optimizer.step = MagicMock(return_value=None)
self.optimizer.zero_grad = MagicMock(return_value=None)
else:
self.optimizer = optimizer
if scheduler_config:
config.training.lr_scheduler = True
config.scheduler = scheduler_config
self.lr_scheduler_callback = LRSchedulerCallback(config, self)
self.callbacks.append(self.lr_scheduler_callback)
on_update_end_fn = (on_update_end_fn if on_update_end_fn else
self.lr_scheduler_callback.on_update_end)
if grad_clipping_config:
self.training_config.clip_gradients = True
self.training_config.max_grad_l2_norm = grad_clipping_config[
"max_grad_l2_norm"]
self.training_config.clip_norm_mode = grad_clipping_config[
"clip_norm_mode"]
self.on_batch_start = MagicMock(return_value=None)
self.on_update_start = MagicMock(return_value=None)
self.logistics_callback = MagicMock(return_value=None)
self.logistics_callback.log_interval = MagicMock(return_value=None)
self.on_batch_end = MagicMock(return_value=None)
self.on_update_end = (on_update_end_fn if on_update_end_fn else
MagicMock(return_value=None))
self.after_training_loop = MagicMock(return_value=None)
self.on_validation_start = MagicMock(return_value=None)
self.scaler = torch.cuda.amp.GradScaler(enabled=False)
self.early_stop_callback = MagicMock(return_value=None)
self.on_validation_end = MagicMock(return_value=None)
self.metrics = MagicMock(return_value={})
self.num_data = num_train_data
def load_datasets(self):
self.dataset_loader = MultiDataModuleNumbersTestObject(
num_data=self.num_data, batch_size=self.config.training.batch_size)
self.dataset_loader.seed_sampler = MagicMock(return_value=None)
self.dataset_loader.prepare_batch = lambda x: SampleList(x)
self.train_loader = self.dataset_loader.train_dataloader()
self.val_loader = self.dataset_loader.val_dataloader()
self.test_loader = self.dataset_loader.test_dataloader()
class TrainerTrainingLoopMock(TrainerTrainingLoopMixin, TrainerProfilingMixin):
def __init__(
self,
num_train_data,
max_updates,
max_epochs,
config=None,
optimizer=None,
update_frequency=1,
batch_size=1,
fp16=False,
on_update_end_fn=None,
):
if config is None:
self.training_config = OmegaConf.create({
"detect_anomaly": False,
"evaluation_interval": 10000,
"update_frequency": update_frequency,
"fp16": fp16,
"batch_size": batch_size,
})
else:
self.training_config = config.training
if max_updates is not None:
self.training_config["max_updates"] = max_updates
if max_epochs is not None:
self.training_config["max_epochs"] = max_epochs
self.model = SimpleModel(1)
if torch.cuda.is_available():
self.model = self.model.cuda()
self.device = "cuda"
else:
self.device = "cpu"
self.dataset_loader = MagicMock()
self.dataset_loader.seed_sampler = MagicMock(return_value=None)
self.dataset_loader.prepare_batch = lambda x: SampleList(x)
if optimizer is None:
self.optimizer = MagicMock()
self.optimizer.step = MagicMock(return_value=None)
self.optimizer.zero_grad = MagicMock(return_value=None)
else:
self.optimizer = optimizer
dataset = NumbersDataset(num_train_data)
self.train_loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=batch_size,
shuffle=False,
num_workers=1,
drop_last=False,
)
self.train_loader.current_dataset = dataset
self.on_batch_start = MagicMock(return_value=None)
self.on_update_start = MagicMock(return_value=None)
self.logistics_callback = MagicMock(return_value=None)
self.logistics_callback.log_interval = MagicMock(return_value=None)
self.on_batch_end = MagicMock(return_value=None)
self.on_update_end = (on_update_end_fn if on_update_end_fn else
MagicMock(return_value=None))
self.meter = Meter()
self.after_training_loop = MagicMock(return_value=None)
self.on_validation_start = MagicMock(return_value=None)
self.evaluation_loop = MagicMock(return_value=(None, None))
self.scaler = torch.cuda.amp.GradScaler(enabled=False)
self.val_loader = MagicMock(return_value=None)
self.early_stop_callback = MagicMock(return_value=None)
self.on_validation_end = MagicMock(return_value=None)
self.metrics = MagicMock(return_value=None)