def instantiate_system():
system = sample_system_object_torch()
system.pipeline.ops = [
fe.op.numpyop.meta.Sometimes(
TestNumpyOp(inputs="x", outputs="x", mode="train", var=1))
]
return system
def test_torch_on_epoch_end(self):
tensorboard = TensorBoard(log_dir=self.log_dir,
weight_histogram_freq=1,
update_freq=1,
write_images='images',
write_embeddings='embed',
embedding_images='embed_images')
tensorboard.system = sample_system_object_torch()
tensorboard.system.global_step = 1
tensorboard.writer = _TorchWriter(self.log_dir, '', tensorboard.system.network)
model = fe.build(model_fn=fe.architecture.pytorch.LeNet, optimizer_fn='adam')
tensorboard.system.network.epoch_models = {model}
if os.path.exists(self.train_path):
shutil.rmtree(self.train_path)
tensorboard.on_epoch_end(data=self.torch_data)
tsv_path = os.path.join(self.embed_path, 'tensors.tsv')
embed_img_path = os.path.join(self.embed_path, 'sprite.png')
# get tensor data from tsv file
fo = open(tsv_path)
tsv_content = csv.reader(fo, delimiter='\t')
for row in tsv_content:
tsv_data = row
fo.close()
# get the image data
output_img = np.asarray(Image.open(embed_img_path))
with self.subTest('Check if tensors.tsv was generated'):
self.assertTrue(os.path.exists(tsv_path))
with self.subTest('Check if embed image was generated'):
self.assertTrue(os.path.exists(embed_img_path))
with self.subTest('Check content of tensors.tsv'):
self.assertEqual(tsv_data, 27 * ['1.0'])
with self.subTest('Check embed image content'):
self.assertTrue(is_equal(output_img, 255 * np.ones(shape=(3, 3, 3), dtype=np.int)))
def instantiate_system():
system = sample_system_object_torch()
model = fe.build(model_fn=fe.architecture.pytorch.LeNet, optimizer_fn='adam', model_name='torch')
system.network = fe.Network(ops=[
fe.op.tensorop.meta.Sometimes(TestTensorOp(inputs="x_out", outputs="x_out", mode="train", var=1)),
ModelOp(model=model, inputs="x_out", outputs="y_pred")
])
return system
def test_torch_on_begin(self):
tensorboard = TensorBoard(log_dir=self.log_dir)
tensorboard.system = sample_system_object_torch()
tensorboard.system.global_step = 1
with patch('sys.stdout', new=StringIO()) as fake_stdout:
tensorboard.on_begin(data=self.torch_data)
log = fake_stdout.getvalue().strip()
self.assertEqual(log, self.on_begin_msg)
def instantiate_system():
system = sample_system_object_torch()
x_train = np.ones((2, 3, 28, 28))
y_train = np.ones((2, ))
ds = TestDataset(data={'x': x_train, 'y': y_train}, var=1)
train_data = fe.dataset.BatchDataset(datasets=[ds, ds], num_samples=[1, 1])
system.pipeline = fe.Pipeline(train_data=train_data, batch_size=2)
return system
def instantiate_system():
system = sample_system_object_torch()
system.pipeline.ops = [
EpochScheduler(epoch_dict={
1:
TestNumpyOp(inputs="x", outputs="x", mode="train", var=1)
})
]
return system
def instantiate_system():
system = sample_system_object_torch()
system.pipeline.ops = [
RepeatScheduler([
TestNumpyOp(inputs="x", outputs="x", mode="train", var=1),
TestNumpyOp(inputs="x", outputs="x", mode="train", var=1)
])
]
return system
def instantiate_system():
system = sample_system_object_torch()
system.pipeline.ops = [
fe.op.numpyop.meta.Repeat(op=TestNumpyOp(inputs="x",
outputs="x",
mode="train",
var=1),
repeat=2)
]
return system
def instantiate_system():
system = sample_system_object_torch()
x_train = np.ones((2, 3, 28, 28))
y_train = np.ones((2, ))
train_data = EpochScheduler(
epoch_dict={
1: TestDataset(data={
'x': x_train,
'y': y_train
}, var=1)
})
system.pipeline = fe.Pipeline(train_data=train_data, batch_size=1)
return system
def test_save_and_load_state_torch(self):
system = sample_system_object_torch()
model_names = get_model_names(system)
global_step = 100
epoch_idx = 10
save_path = tempfile.mkdtemp()
system.global_step = global_step
system.epoch_idx = epoch_idx
with self.subTest("Check state files were created"):
system.save_state(save_dir=save_path)
self.assertTrue(os.path.exists(os.path.join(save_path, 'ds.pkl')))
self.assertTrue(os.path.exists(os.path.join(save_path,
'nops.pkl')))
self.assertTrue(
os.path.exists(os.path.join(save_path, 'summary.pkl')))
self.assertTrue(
os.path.exists(os.path.join(save_path, 'system.json')))
self.assertTrue(os.path.exists(os.path.join(save_path,
'tops.pkl')))
self.assertTrue(
os.path.exists(os.path.join(save_path, 'traces.pkl')))
for model_name in model_names:
self.assertTrue(
os.path.exists(os.path.join(save_path,
f'{model_name}.pt')))
self.assertTrue(
os.path.exists(
os.path.join(save_path, f'{model_name}_opt.pt')))
system = sample_system_object_torch()
with self.subTest("Check that state loads properly"):
system.load_state(save_path)
self.assertEqual(system.global_step, global_step)
self.assertEqual(system.epoch_idx, epoch_idx)
if os.path.exists(save_path):
shutil.rmtree(save_path)
def instantiate_system():
system = sample_system_object_torch()
model = fe.build(model_fn=fe.architecture.pytorch.LeNet,
optimizer_fn='adam',
model_name='torch')
var1 = torch.tensor(1.0)
system.network = fe.Network(ops=[
TestTensorOp(
inputs="x_out", outputs="x_out", mode="train", var1=var1),
ModelOp(model=model, inputs="x_out", outputs="y_pred")
])
system.traces.append(TestTrace(var1=var1))
return system
def test_save_and_load_state_torch(self):
""" `save_state` and `load_state` of an entire system are highly dependent on the implementation of __setstate__
and __getstate__ of ops, traces, datasets ... etc. The save_state and load_state function should be tested in
the testing script of that perticular component
(ex: The test of save_state and load_state with EpochScheduler is located in
PR_test/integration_test/schedule/test_epoch_scheduler.py)
"""
system = sample_system_object_torch()
model_names = get_model_names(system)
global_step = 100
epoch_idx = 10
save_path = tempfile.mkdtemp()
system.global_step = global_step
system.epoch_idx = epoch_idx
with self.subTest("Check state files were created"):
system.save_state(save_dir=save_path)
self.assertTrue(
os.path.exists(os.path.join(save_path, 'objects.pkl')))
self.assertTrue(
os.path.exists(os.path.join(save_path, 'system.json')))
for model_name in model_names:
self.assertTrue(
os.path.exists(os.path.join(save_path,
f'{model_name}.pt')))
self.assertTrue(
os.path.exists(
os.path.join(save_path, f'{model_name}_opt.pt')))
system = sample_system_object_torch()
with self.subTest("Check that state loads properly"):
system.load_state(save_path)
self.assertEqual(global_step, system.global_step)
self.assertEqual(epoch_idx, system.epoch_idx)
if os.path.exists(save_path):
shutil.rmtree(save_path)
def instantiate_system():
system = sample_system_object_torch()
submodel = OneLayerTorchModel()
model = fe.build(model_fn=lambda: TestModel(submodel),
optimizer_fn='adam',
model_name='torch')
model2 = fe.build(model_fn=lambda: TestModel(submodel),
optimizer_fn='adam',
model_name='torch2')
system.network = fe.Network(ops=[
ModelOp(model=model, inputs="x_out", outputs="y_pred"),
ModelOp(model=model2, inputs="x_out", outputs="y_pred2"),
])
return system
def test_torch_model_on_begin(self):
restore_wizard = RestoreWizard(directory=self.system_json_path)
restore_wizard.system = sample_system_object_torch()
# save state
for model in restore_wizard.system.network.models:
save_model(model,
save_dir=restore_wizard.directory,
save_optimizer=True)
restore_wizard.system.save_state(json_path=os.path.join(
restore_wizard.directory, restore_wizard.system_file))
restore_wizard.on_begin(data=self.data)
with self.subTest('Check the restore files directory'):
self.assertEqual(restore_wizard.directory, self.system_json_path)
with self.subTest('check data dictionary'):
self.assertEqual(self.data['epoch'], 0)
if os.path.exists(self.system_json_path):
shutil.rmtree(self.system_json_path)
def instantiate_system():
system = sample_system_object_torch()
model = fe.build(model_fn=fe.architecture.pytorch.LeNet,
optimizer_fn='adam',
model_name='torch')
system.network = fe.Network(ops=[
EpochScheduler(
epoch_dict={
1:
TestTensorOp(inputs="x_out",
outputs="x_out",
mode="train",
var=1)
}),
ModelOp(model=model, inputs="x_out", outputs="y_pred")
])
return system
def test_torch_on_batch_end(self):
tensorboard = TensorBoard(log_dir=self.log_dir, weight_histogram_freq=1, update_freq=1)
tensorboard.system = sample_system_object_torch()
tensorboard.system.global_step = 1
tensorboard.writer = _TorchWriter(self.log_dir, '', tensorboard.system.network)
model = fe.build(model_fn=fe.architecture.pytorch.LeNet, optimizer_fn='adam', model_name='torch')
model.fe_input_spec = FeInputSpec(self.torch_data['x'], model)
tensorboard.system.network.epoch_models = {model}
if os.path.exists(self.train_path):
shutil.rmtree(self.train_path)
tensorboard.on_batch_end(data=self.torch_data)
filepath = getfilepath()
for e in tf.compat.v1.train.summary_iterator(filepath):
for v in e.summary.value:
if v.tag == "torch_fc1/bias":
output = v.histo.num
self.assertEqual(output, 64.0)
def instantiate_system():
system = sample_system_object_torch()
x_train = np.ones((2, 3, 28, 28))
y_train = np.ones((2, ))
train_data = RepeatScheduler([
TestDataset(data={
"x": x_train,
"y": y_train
}, var=1),
TestDataset(data={
"x": x_train,
"y": y_train
}, var=1),
])
system.pipeline = fe.Pipeline(train_data=train_data, batch_size=1)
return system
def test_torch_model_on_epoch_end(self):
restore_wizard = RestoreWizard(directory=self.system_json_path)
restore_wizard.system = sample_system_object_torch()
restore_wizard.on_epoch_end(data=self.data)
model_names = get_model_name(restore_wizard.system)
with self.subTest('check json exists'):
self.assertTrue(
os.path.exists(
os.path.join(self.system_json_path, 'system.json')))
with self.subTest('Check if model weights path stored'):
self.assertTrue(
os.path.exists(
os.path.join(self.system_json_path,
model_names[0] + '.pt')))
with self.subTest('Check if model optimizer stored'):
self.assertTrue(
os.path.exists(
os.path.join(self.system_json_path,
model_names[0] + '_opt.pt')))
if os.path.exists(self.system_json_path):
shutil.rmtree(self.system_json_path)
def instantiate_system():
system = sample_system_object_torch()
system.traces.append(
EpochScheduler(epoch_dict={1: TestTrace(var1=1)}))
return system
def instantiate_system():
system = sample_system_object_torch()
system.traces.append(
RepeatScheduler([TestTrace(var1=1),
TestTrace(var1=1)]))
return system
