def test_trainer_invalid_options(path):
device = 'cpu'
model = MyModel()
ppe.to(model, device)
model_with_loss = MyModelWithLossFn(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
extensions = _make_extensions()
options = {'UNKNOWN_OPTIONS': True}
with pytest.raises(ValueError, match="UNKNOWN_OPTIONS"):
engine.create_trainer(
model_with_loss, optimizer, 20,
device=device, extensions=extensions,
out_dir=path,
options=options,
)
def test_trainer_with_code_block(device, progress_bar, path):
model = MyModel()
model_with_loss = MyModelWithLossDictOutput(model)
ppe.to(model_with_loss, device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader([{
'x': torch.rand(20, ),
't': torch.rand(10, )
} for i in range(10)])
extensions = _make_extensions()
evaluator = engine.create_evaluator(model_with_loss,
device=device,
progress_bar=progress_bar,
logic=ppe.handler.CodeBlockLogic())
trainer = engine.create_trainer(model_with_loss,
optimizer,
20,
device=device,
evaluator=evaluator,
extensions=extensions,
out_dir=path,
logic=ppe.handler.CodeBlockLogic())
trainer.run(data, data)
def test_trainer_defer_wrong_order(path):
class WrongOrderHandler(ppe.handler.Handler):
def _complete_train_step(self, trainer, outs, block, sn, sm, rt):
p_iter = self.pending_iters[sn][0]
if p_iter.idx < 10:
super()._complete_train_step(trainer, p_iter.deferred, block,
sn, sm, rt)
else:
p_iter.cback(90, None, is_deferred=block)
device = 'cpu'
model = MyModel()
model_with_loss = MyModelWithLossAsync(model)
ppe.to(model_with_loss, device)
# Register the handler
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader([(torch.rand(20, ), torch.rand(10, ))
for i in range(100)])
trainer = engine.create_trainer(model_with_loss,
optimizer,
2,
device=device,
handler_class=WrongOrderHandler,
out_dir=path)
with pytest.raises(RuntimeError, match="Completed a not expected"):
trainer.run(data)
def test_trainer_defer(path):
class Extension:
def __init__(self, is_async):
self.name = 'Dummy'
self.trigger = (1, 'iteration')
self.called = 0
self.is_async = is_async
def __call__(self, manager):
self.called += 1
device = 'cpu'
model = MyModel()
model_with_loss = MyModelWithLossAsync(model)
ppe.to(model_with_loss, device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader([(torch.rand(20, ), torch.rand(10, ))
for i in range(100)])
extensions = [Extension(True), Extension(False)]
trainer = engine.create_trainer(model_with_loss,
optimizer,
2,
device=device,
extensions=extensions,
out_dir=path)
trainer.run(data)
assert trainer.manager.iteration == 200
assert trainer.manager.execution == 200
assert extensions[0].called == 200
assert extensions[1].called == 200
def test_evaluator_trigger(evaluator_trigger, path):
device = 'cpu'
progress_bar = False
model = MyModel()
ppe.to(model, device)
model_with_loss = MyModelWithLossFn(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader([(torch.rand(20, ), torch.rand(10, ))
for i in range(10)])
extensions = _make_extensions()
evaluator = engine.create_evaluator(model_with_loss,
device=device,
progress_bar=progress_bar)
trainer = engine.create_trainer(model_with_loss,
optimizer,
20,
device=device,
evaluator=(evaluator,
evaluator_trigger[1]),
extensions=extensions,
out_dir=path)
path = 'pytorch_pfn_extras.training._evaluator.Evaluator.run'
with mock.patch(path) as patched:
trainer.run(data, data)
assert patched.call_count == evaluator_trigger[0]
def test_trainer_profile():
device = 'cpu'
model = MyModel()
model_with_loss = MyModelWithLossDictOutput(model)
ppe.to(model_with_loss, device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader([{
'x': torch.rand(20, ),
't': torch.rand(10, )
} for i in range(10)])
extensions = _make_extensions()
evaluator = engine.create_evaluator(model_with_loss, device=device)
trace_handler = mock.Mock()
warmup = 1
active = len(data) - warmup
profile = torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU],
on_trace_ready=trace_handler,
schedule=torch.profiler.schedule(wait=0, warmup=warmup, active=active),
)
trainer = engine.create_trainer(
model_with_loss,
optimizer,
20,
device=device,
evaluator=evaluator,
extensions=extensions,
profile=profile,
)
trainer.run(data, data)
assert trace_handler.call_count == 20 # n_epochs
def test_trainer_with_code_block_with_multiple_optimizers(
device, progress_bar, path):
if not torch.cuda.is_available() and device == 'cuda':
pytest.skip()
model = MyModel()
model_with_loss = MyModelWithLossDictOutput(model)
ppe.to(model_with_loss, device)
optimizer0 = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer1 = torch.optim.Adam(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader([{
'x': torch.rand(20, ),
't': torch.rand(10, )
} for i in range(10)])
extensions = _make_extensions()
evaluator = engine.create_evaluator(model_with_loss,
device=device,
progress_bar=progress_bar,
logic=ppe.handler.CodeBlockLogic())
trainer = engine.create_trainer(model_with_loss, {
"0": optimizer0,
"1": optimizer1
},
20,
device=device,
evaluator=evaluator,
extensions=extensions,
out_dir=path,
logic=ppe.handler.CodeBlockLogic())
trainer.run(data, data)
def test_trainer_dict_input(device, progress_bar, path):
if not torch.cuda.is_available() and device == 'cuda':
pytest.skip()
model = MyModel()
ppe.to(model, device)
model_with_loss = MyModelWithLossDictOutput(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader([{
'x': torch.rand(20, ),
't': torch.rand(10, )
} for i in range(10)])
extensions = _make_extensions()
evaluator = engine.create_evaluator(model_with_loss,
device=device,
progress_bar=progress_bar)
trainer = engine.create_trainer(model_with_loss,
optimizer,
20,
device=device,
evaluator=evaluator,
extensions=extensions,
out_dir=path)
trainer.run(data, data)
def test_train_with_evaluator(device, progress_bar, path):
if not torch.cuda.is_available() and device == 'cuda':
pytest.skip()
model = MyModel()
ppe.to(model, device)
model_with_loss = MyModelWithLossFn(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader([(torch.rand(20, ), torch.rand(10, ))
for i in range(10)])
extensions = _make_extensions()
evaluator = engine.create_evaluator(model_with_loss,
device=device,
progress_bar=progress_bar)
trainer = engine.create_trainer(model_with_loss,
optimizer,
20,
device=device,
evaluator=evaluator,
extensions=extensions,
out_dir=path)
mpath = 'pytorch_pfn_extras.training._evaluator.Evaluator.run'
with mock.patch(mpath) as patched:
trainer.run(data, data)
assert patched.call_count == 20
def _get_trainer(self, epochs, out_dir):
model = MyModel()
ppe.to(model, 'cpu')
model_with_loss = MyModelWithLossFn(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
extensions = _make_extensions()
trainer = engine.create_trainer(
model_with_loss, optimizer, 20,
device='cpu', extensions=extensions,
out_dir=out_dir
)
return trainer
def test_trainer_no_to(path):
model = MyModel()
model_with_loss = MyModelWithLossFn(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader(
[(torch.rand(20,), torch.rand(10,)) for i in range(10)])
extensions = _make_extensions()
trainer = engine.create_trainer(
model_with_loss, optimizer, 20,
device='cpu', extensions=extensions,
out_dir=path,
)
with pytest.raises(RuntimeError, match="ppe.to"):
trainer.run(data, data)
def test_trainer(device, path):
if not torch.cuda.is_available() and device == 'cuda':
pytest.skip()
model = MyModel()
ppe.to(model, device)
model_with_loss = MyModelWithLossFn(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader(
[(torch.rand(20,), torch.rand(10,)) for i in range(10)])
extensions = _make_extensions()
trainer = engine.create_trainer(
model_with_loss, optimizer, 20,
device=device, extensions=extensions,
out_dir=path,
)
trainer.run(data)
def get_result_from_trainer():
model = MyModel()
ppe.to(model, device)
model_with_loss = MyModelWithLossFn(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
extensions = _make_extensions()
trainer = engine.create_trainer(
model_with_loss, optimizer, 20,
device=device, extensions=extensions,
out_dir=path
)
trainer.run(train_data)
model.eval()
with torch.no_grad():
return [model(x.to(device)) for x, in data]
def test_trainer_namedtuple_input(device, progress_bar, path):
model = MyModel()
ppe.to(model, device)
model_with_loss = ModelNamedTupleIO(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader(
[Input(torch.rand(20, ), torch.rand(10, ), str(i)) for i in range(10)])
extensions = _make_extensions()
evaluator = engine.create_evaluator(model_with_loss,
device=device,
progress_bar=progress_bar)
trainer = engine.create_trainer(model_with_loss,
optimizer,
20,
device=device,
evaluator=evaluator,
extensions=extensions,
out_dir=path)
trainer.run(data, data)
def add_dump(self, name, dir):
with open(f'{dir}/summary') as f:
summary = json.loads(f.read())
class DummyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.param = torch.nn.Parameter(torch.tensor(1.0))
def forward(self, *args, **kwargs):
return ()
model = DummyModel()
ppe.to(model, 'cpu')
engine = None
args = []
if summary['evaluator']:
engine = _engine_module.create_evaluator(model)
args = [[None] * summary['eval_len']]
if summary['trainer']:
engine = _engine_module.create_trainer(
{'main': model},
{'main': torch.optim.SGD(model.parameters(), lr=0.01)},
summary['max_epochs'],
evaluator=engine,
)
args = [[None] * summary['train_len']] + args
engine.handler = _ComparableHandler(
engine.handler, name, self._compare_dump, self._trigger, dir=dir)
child_evaluator = getattr(engine, 'evaluator', None)
if child_evaluator is not None:
# For trainer with evaluator
child_evaluator.handler = engine.handler
self._engines[name] = engine, args, {}
