def test_optimizers_and_schedulers(tmpdir, optim, sched, interval):
model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10),
nn.LogSoftmax())
task = ClassificationTask(model, optimizer=optim, lr_scheduler=sched)
train_dl = torch.utils.data.DataLoader(DummyDataset())
if sched is None:
optimizer = task.configure_optimizers()
assert isinstance(optimizer, torch.optim.Adadelta)
else:
optimizer, scheduler = task.configure_optimizers()
assert isinstance(optimizer[0], torch.optim.Adadelta)
scheduler = scheduler[0]
assert isinstance(scheduler["scheduler"],
torch.optim.lr_scheduler.StepLR)
assert scheduler["interval"] == interval
# generate a checkpoint
trainer = flash.Trainer(
default_root_dir=tmpdir,
limit_train_batches=10,
max_epochs=1,
)
trainer.fit(task, train_dl)
def test_optimization(tmpdir):
model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.LogSoftmax())
optim = torch.optim.Adam(model.parameters())
task = ClassificationTask(model, optimizer=optim, scheduler=None)
optimizer = task.configure_optimizers()
assert optimizer == optim
task = ClassificationTask(model, optimizer=torch.optim.Adadelta, optimizer_kwargs={"eps": 0.5}, scheduler=None)
optimizer = task.configure_optimizers()
assert isinstance(optimizer, torch.optim.Adadelta)
assert optimizer.defaults["eps"] == 0.5
task = ClassificationTask(
model,
optimizer=torch.optim.Adadelta,
scheduler=torch.optim.lr_scheduler.StepLR,
scheduler_kwargs={"step_size": 1}
)
optimizer, scheduler = task.configure_optimizers()
assert isinstance(optimizer[0], torch.optim.Adadelta)
assert isinstance(scheduler[0], torch.optim.lr_scheduler.StepLR)
optim = torch.optim.Adadelta(model.parameters())
task = ClassificationTask(model, optimizer=optim, scheduler=torch.optim.lr_scheduler.StepLR(optim, step_size=1))
optimizer, scheduler = task.configure_optimizers()
assert isinstance(optimizer[0], torch.optim.Adadelta)
assert isinstance(scheduler[0], torch.optim.lr_scheduler.StepLR)
if _TRANSFORMERS_AVAILABLE:
from transformers.optimization import get_linear_schedule_with_warmup
assert task.available_schedulers() == [
'constant_schedule', 'constant_schedule_with_warmup', 'cosine_schedule_with_warmup',
'cosine_with_hard_restarts_schedule_with_warmup', 'linear_schedule_with_warmup',
'polynomial_decay_schedule_with_warmup'
]
optim = torch.optim.Adadelta(model.parameters())
with pytest.raises(MisconfigurationException, match="The LightningModule isn't attached to the trainer yet."):
task = ClassificationTask(model, optimizer=optim, scheduler="linear_schedule_with_warmup")
optimizer, scheduler = task.configure_optimizers()
task = ClassificationTask(
model,
optimizer=optim,
scheduler="linear_schedule_with_warmup",
scheduler_kwargs={"num_warmup_steps": 0.1},
loss_fn=F.nll_loss,
)
trainer = flash.Trainer(max_epochs=1, limit_train_batches=2)
ds = DummyDataset()
trainer.fit(task, train_dataloader=DataLoader(ds))
optimizer, scheduler = task.configure_optimizers()
assert isinstance(optimizer[0], torch.optim.Adadelta)
assert isinstance(scheduler[0], torch.optim.lr_scheduler.LambdaLR)
expected = get_linear_schedule_with_warmup.__name__
assert scheduler[0].lr_lambdas[0].__qualname__.split('.')[0] == expected
def test_external_optimizers_torch_optimizer(tmpdir, optim):
model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10),
nn.LogSoftmax())
task = ClassificationTask(model,
optimizer=optim,
lr_scheduler=None,
loss_fn=F.nll_loss)
trainer = flash.Trainer(max_epochs=1,
limit_train_batches=2,
gpus=torch.cuda.device_count())
ds = DummyDataset()
trainer.fit(task, train_dataloader=DataLoader(ds))
from torch_optimizer import Yogi
optimizer = task.configure_optimizers()
assert isinstance(optimizer, Yogi)
def test_errors_and_exceptions_optimizers_and_schedulers():
model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10),
nn.LogSoftmax())
with pytest.raises(TypeError):
task = ClassificationTask(model,
optimizer=[1, 2, 3, 4],
lr_scheduler=None)
task.configure_optimizers()
with pytest.raises(KeyError):
task = ClassificationTask(model,
optimizer="not_a_valid_key",
lr_scheduler=None)
task.configure_optimizers()
with pytest.raises(TypeError):
task = ClassificationTask(model,
optimizer=(["not", "a", "valid", "type"], {
"random_kwarg": 10
}),
lr_scheduler=None)
task.configure_optimizers()
with pytest.raises(TypeError):
task = ClassificationTask(model,
optimizer=("Adam", ["non", "dict", "type"]),
lr_scheduler=None)
task.configure_optimizers()
with pytest.raises(KeyError):
task = ClassificationTask(model,
optimizer="Adam",
lr_scheduler="not_a_valid_key")
task.configure_optimizers()
with pytest.raises(TypeError):
task = ClassificationTask(model,
optimizer="Adam",
lr_scheduler=["not", "a", "valid", "type"])
task.configure_optimizers()
with pytest.raises(TypeError):
task = ClassificationTask(model,
optimizer="Adam",
lr_scheduler=(["not", "a", "valid",
"type"], {
"random_kwarg": 10
}))
task.configure_optimizers()
