def test_tf_on_begin(self):
tensorboard = TensorBoard(log_dir=self.log_dir)
tensorboard.system = sample_system_object()
tensorboard.system.global_step = 1
with patch('sys.stdout', new=StringIO()) as fake_stdout:
tensorboard.on_begin(data=self.tf_data)
log = fake_stdout.getvalue().strip()
self.assertEqual(log, self.on_begin_msg)
def get_estimator(epochs=2,
batch_size=32,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1
train_data, eval_data = mnist.load_data()
test_data = eval_data.split(0.5)
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[
ExpandDims(inputs="x", outputs="x"),
Minmax(inputs="x", outputs="x")
],
num_process=0)
# step 2
model = fe.build(model_fn=LeNet, optimizer_fn="adam")
print([f"{idx}: {x.name}" for idx, x in enumerate(model.submodules)])
network = fe.Network(ops=[
Watch(inputs="x"),
ModelOp(model=model,
inputs="x",
outputs=["y_pred", "embedding"],
intermediate_layers='dense'),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
GradientOp(finals="embedding", inputs="x", outputs="grads"),
UpdateOp(model=model, loss_name="ce")
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
Inspector(),
BestModelSaver(model=model,
save_dir=save_dir,
metric="accuracy",
save_best_mode="max"),
LRScheduler(model=model,
lr_fn=lambda step: cosine_decay(
step, cycle_length=3750, init_lr=1e-3)),
TensorBoard(log_dir="tf_logs",
write_embeddings="embedding",
embedding_labels="y")
]
estimator = fe.Estimator(
pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch)
return estimator
def test_tf_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()
tensorboard.system.global_step = 1
tensorboard.writer = _TfWriter(self.log_dir, '', tensorboard.system.network)
model = fe.build(model_fn=fe.architecture.tensorflow.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.tf_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 test_tf_on_batch_end(self):
tensorboard = TensorBoard(log_dir=self.log_dir, weight_histogram_freq=1, update_freq=1)
tensorboard.system = sample_system_object()
tensorboard.system.global_step = 1
tensorboard.writer = _TfWriter(self.log_dir, '', tensorboard.system.network)
model = fe.build(model_fn=fe.architecture.tensorflow.LeNet, optimizer_fn='adam')
tensorboard.system.network.epoch_models = {model}
if os.path.exists(self.train_path):
shutil.rmtree(self.train_path)
tensorboard.on_batch_end(data=self.tf_data)
filepath = getfilepath()
for e in tf.compat.v1.train.summary_iterator(filepath):
for v in e.summary.value:
if v.tag == "tf_dense_1/bias_0":
output = v.histo.num
self.assertEqual(output, 10.0)
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)