def test_checkpoint_save_and_restore_via_public_trainer_methods(
dev_str, call, compile_mode):
if call is helpers.np_call:
# numpy does not support gradients, required for training
pytest.skip()
# currently only PyTorch supports graph compilation
compile_mode = compile_mode if ivy.current_framework_str(
) == 'torch' else False
builder_helpers.remove_dirs()
data_loader_spec_args = {'batch_size': 1, 'dev_strs': [dev_str]}
trainer_spec_args = {
'total_iterations': 0,
'ld_chkpt': False,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoaderMin,
ExampleNetworkMin,
ExampleTrainerMin,
data_loader_spec_args=data_loader_spec_args,
trainer_spec_args=trainer_spec_args)
trainer.setup()
chkpt0_path = os.path.join('chkpt/', 'test_chkpt0.hdf5')
trainer.save(chkpt0_path)
assert os.path.exists(chkpt0_path)
trainer.train()
chkpt1_path = os.path.join('chkpt/', 'test_chkpt1.hdf5')
trainer.save(chkpt1_path)
trainer.close()
assert os.path.exists(chkpt1_path)
data_loader_spec_args = {'batch_size': 1, 'dev_strs': [dev_str]}
trainer_spec_args = {
'total_iterations': 10,
'ld_chkpt': False,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoaderMin,
ExampleNetworkMin,
ExampleTrainerMin,
data_loader_spec_args=data_loader_spec_args,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.restore(chkpt0_path)
trainer.train()
chkpt3_path = os.path.join('chkpt/', 'test_chkpt3.hdf5')
trainer.save(chkpt3_path)
trainer.close()
assert os.path.exists(chkpt3_path)
builder_helpers.remove_dirs()
def test_visualizing(dev_str, call, compile_mode):
if call is helpers.np_call:
# numpy does not support gradients, required for training
pytest.skip()
# currently only PyTorch supports graph compilation
compile_mode = compile_mode if ivy.current_framework_str(
) == 'torch' else False
builder_helpers.remove_dirs()
data_loader_spec_args = {'batch_size': 1, 'dev_strs': [dev_str]}
trainer_spec_args = {
'total_iterations': 10,
'ld_chkpt': False,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoaderMin,
ExampleNetworkMin,
ExampleTrainerMin,
data_loader_spec_args=data_loader_spec_args,
trainer_spec_args=trainer_spec_args)
trainer.setup()
try:
trainer.visualize()
except OSError:
pass
trainer.close()
builder_helpers.remove_dirs()
def test_checkpoint_loading(dev_str, call, compile_mode):
if call is helpers.np_call:
# numpy does not support gradients, required for training
pytest.skip()
# currently only PyTorch supports graph compilation
compile_mode = compile_mode if ivy.current_framework_str(
) == 'torch' else False
builder_helpers.remove_dirs()
data_loader_spec_args = {'batch_size': 1, 'dev_strs': [dev_str]}
trainer_spec_args = {
'total_iterations': 10,
'ld_chkpt': False,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoaderMin,
ExampleNetworkMin,
ExampleTrainerMin,
data_loader_spec_args=data_loader_spec_args,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
trainer.close()
trainer_spec_args = {
'total_iterations': 20,
'ld_chkpt': True,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoaderMin,
ExampleNetworkMin,
ExampleTrainerMin,
data_loader_spec_args=data_loader_spec_args,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
trainer.close()
checkpoint_nums = [
int(fname.split('-')[-1].split('.')[0])
for fname in os.listdir('log/chkpts')
]
assert max(checkpoint_nums) == 19
builder_helpers.remove_dirs()
def main(compile_mode=False):
current_dir = os.path.dirname(os.path.abspath(__file__))
# dataset dirs specification
dataset_dirs_args = dict()
# dataset specification
dataset_spec_filepath = os.path.join(current_dir, 'json_specs',
'dataset_spec.json.example')
dataset_spec_args = builder.parse_json_to_cont(dataset_spec_filepath)
# data loader specification
data_loader_spec_filepath = os.path.join(current_dir, 'json_specs',
'data_loader_spec.json.example')
data_loader_spec_args = builder.parse_json_to_cont(
data_loader_spec_filepath)
# network specification
network_spec_filepath = os.path.join(current_dir, 'json_specs',
'network_spec.json.example')
network_spec_args = builder.parse_json_to_cont(network_spec_filepath)
# trainer specification
trainer_spec_filepath = os.path.join(current_dir, 'json_specs',
'trainer_spec.json.example')
trainer_spec_args = builder.parse_json_to_cont(trainer_spec_filepath)
# In all above cases, the user could override the loaded json file dicts with command line args if so desired
# before then passing into the TrainingJob for specification class construction, which are all then read-only
trainer = builder.build_trainer(
ExampleDataLoader,
ExampleNetwork,
ExampleTrainer,
dataset_dirs_args=dataset_dirs_args,
dataset_dirs_class=ExampleDatasetDirs,
dataset_spec_args=dataset_spec_args,
dataset_spec_class=ExampleDatasetSpec,
data_loader_spec_args=data_loader_spec_args,
data_loader_spec_class=ExampleDataLoaderSpec,
network_spec_args=network_spec_args,
network_spec_class=ExampleNetworkSpec,
trainer_spec_args=trainer_spec_args,
spec_cont=ivy.Container({'trainer': {
'compile_mode': compile_mode
}}))
trainer.setup()
print("Finished complete example!")
trainer.train()
trainer.close()
def main(seed=0, compile_mode=False, dev_strs=None):
ivy.seed(seed)
data_loader_spec_args = {
'batch_size': 2,
'dev_strs': [ivy.default(lambda: dev_strs[0], None, True)]
}
trainer_spec_args = {
'total_iterations': 10,
'ld_chkpt': False,
'log_freq': 1,
'initial_learning_rate': 0.1,
'compile_mode': compile_mode,
'dev_strs': dev_strs
}
trainer = trainer_builder.build_trainer(
ExampleDataLoader,
ExampleNetwork,
ExampleTrainer,
data_loader_spec_args=data_loader_spec_args,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
trainer.close()
def test_reduced_cost_after_checkpoint_load(dev_str, call, compile_mode):
if call is helpers.np_call:
# numpy does not support gradients, required for training
pytest.skip()
if call is helpers.jnp_call and ivy.wrapped_mode():
# Jax does not support ivy.Array instances when calling _jax.grad()
pytest.skip()
example_dir = os.path.relpath(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'../ivy_builder_demos'))
# currently only PyTorch supports graph compilation
compile_mode = compile_mode if ivy.current_framework_str(
) == 'torch' else False
# dataset dirs specification
dataset_dirs_args = dict()
# dataset specification
dataset_spec_filepath = os.path.join(example_dir, 'json_specs',
'dataset_spec.json.example')
dataset_spec_args = builder.parse_json_to_cont(dataset_spec_filepath)
# data loader specification
data_loader_spec_filepath = os.path.join(example_dir, 'json_specs',
'data_loader_spec.json.example')
data_loader_spec_args = builder.parse_json_to_cont(
data_loader_spec_filepath)
# network specification
network_spec_filepath = os.path.join(example_dir, 'json_specs',
'network_spec.json.example')
network_spec_args = builder.parse_json_to_cont(network_spec_filepath)
builder_helpers.remove_dirs()
ivy.seed(0)
trainer_spec_args = {
'total_iterations': 1,
'ld_chkpt': False,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoader,
ExampleNetwork,
ExampleTrainer,
dataset_dirs_args=dataset_dirs_args,
dataset_dirs_class=ExampleDatasetDirs,
dataset_spec_args=dataset_spec_args,
dataset_spec_class=ExampleDatasetSpec,
data_loader_spec_args=data_loader_spec_args,
data_loader_spec_class=ExampleDataLoaderSpec,
network_spec_args=network_spec_args,
network_spec_class=ExampleNetworkSpec,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
initial_cost = trainer._total_cost
assert trainer._global_step == 1
trainer.close()
ivy.seed(0)
steps_to_take_first = 10
trainer_spec_args = {
'total_iterations': steps_to_take_first,
'ld_chkpt': False,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoader,
ExampleNetwork,
ExampleTrainer,
dataset_dirs_args=dataset_dirs_args,
dataset_dirs_class=ExampleDatasetDirs,
dataset_spec_args=dataset_spec_args,
dataset_spec_class=ExampleDatasetSpec,
data_loader_spec_args=data_loader_spec_args,
data_loader_spec_class=ExampleDataLoaderSpec,
network_spec_args=network_spec_args,
network_spec_class=ExampleNetworkSpec,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
ten_step_cost = trainer._total_cost
assert trainer._global_step == steps_to_take_first
trainer.close()
assert initial_cost > ten_step_cost
steps_to_take_second = 20
trainer_spec_args = {
'total_iterations': steps_to_take_second,
'ld_chkpt': True,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoader,
ExampleNetwork,
ExampleTrainer,
dataset_dirs_args=dataset_dirs_args,
dataset_dirs_class=ExampleDatasetDirs,
dataset_spec_args=dataset_spec_args,
dataset_spec_class=ExampleDatasetSpec,
data_loader_spec_args=data_loader_spec_args,
data_loader_spec_class=ExampleDataLoaderSpec,
network_spec_args=network_spec_args,
network_spec_class=ExampleNetworkSpec,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
twenty_step_cost = trainer._total_cost
assert trainer._global_step == steps_to_take_second
trainer.close()
assert ten_step_cost > twenty_step_cost
builder_helpers.remove_dirs()