def load(self, config, dataset_type="train", *args, **kwargs):
dataset = NumbersDataset(DATASET_LEN,
data_item_key="input",
always_one=True)
dataset.dataset_name = self.dataset_name
dataset.dataset_type = dataset_type
return dataset
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)
def _get_dataloader(self):
dataset = NumbersDataset(self._num_data)
dataloader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=self._batch_size,
shuffle=False,
num_workers=1,
drop_last=False,
)
return dataloader
def _get_dataloader(self, dataset_type="train"):
dataset = NumbersDataset(
self._num_data, always_one=True, dataset_type=dataset_type
)
return torch.utils.data.DataLoader(
dataset=dataset,
batch_size=self.batch_size,
shuffle=False,
num_workers=1,
drop_last=False,
)
def setUp(self):
self.trainer = argparse.Namespace()
self.config = load_yaml(os.path.join("configs", "defaults.yaml"))
self.config = OmegaConf.merge(
self.config,
{
"model": "simple",
"model_config": {},
"training": {
"lr_scheduler": True,
"lr_ratio": 0.1,
"lr_steps": [1, 2],
"use_warmup": False,
"callbacks": [{
"type": "test_callback",
"params": {}
}],
},
},
)
# Keep original copy for testing purposes
self.trainer.config = deepcopy(self.config)
registry.register("config", self.trainer.config)
model = SimpleModel(SimpleModel.Config())
model.build()
self.trainer.model = model
self.trainer.val_loader = torch.utils.data.DataLoader(
NumbersDataset(2), batch_size=self.config.training.batch_size)
self.trainer.optimizer = torch.optim.Adam(
self.trainer.model.parameters(), lr=1e-01)
self.trainer.lr_scheduler_callback = LRSchedulerCallback(
self.config, self.trainer)
self.trainer.callbacks = []
for callback in self.config.training.get("callbacks", []):
callback_type = callback.type
callback_param = callback.params
callback_cls = registry.get_callback_class(callback_type)
self.trainer.callbacks.append(
callback_cls(self.trainer.config, self.trainer,
**callback_param))
def __init__(
self,
config,
max_steps,
max_epochs=None,
callback=None,
num_data_size=100,
batch_size=1,
accumulate_grad_batches=1,
lr_scheduler=False,
gradient_clip_val=0.0,
precision=32,
):
self.config = config
self._callbacks = None
if callback:
self._callbacks = [callback]
# data
self.data_module = MagicMock()
dataset = NumbersDataset(num_data_size)
self.data_module.train_loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=batch_size,
shuffle=False,
num_workers=1,
drop_last=False,
)
self.data_module.train_loader.current_dataset = MagicMock(
return_value=dataset)
# settings
trainer_config = self.config.trainer.params
trainer_config.accumulate_grad_batches = accumulate_grad_batches
trainer_config.precision = precision
trainer_config.max_steps = max_steps
trainer_config.max_epochs = max_epochs
trainer_config.gradient_clip_val = gradient_clip_val
trainer_config.precision = precision
self.trainer_config = trainer_config
self.config.training.lr_scheduler = lr_scheduler
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)
def _build_dataloaders(self):
dataloaders = {}
for idx in range(self.NUM_DATALOADERS):
dataloaders[f"numbers_{idx}"] = torch.utils.data.DataLoader(
dataset=NumbersDataset((idx + 1) * (10**idx)), num_workers=0)
return dataloaders
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)