def test_simple_input(self):
inp = torch.ones(size=(32, 1, 28, 28), dtype=torch.float16)
model = torch.nn.Sequential(torch.nn.Linear(28 * 28, 100))
spec = FeInputSpec(inp, model)
result = spec.get_dummy_input()
self.assertIsInstance(result, torch.Tensor,
"Spec should return a torch tensor")
self.assertListEqual([32, 1, 28, 28], list(result.shape),
"Result shape is incorrect")
self.assertEqual(torch.float16, result.dtype,
"Result should be float16 dtype")
def test_set_input(self):
inp = {
torch.ones(size=(32, 3, 32, 32), dtype=torch.float64),
torch.ones(size=(32, 10), dtype=torch.int8)
}
model = torch.nn.Sequential(torch.nn.Linear(28 * 28, 100))
spec = FeInputSpec(inp, model)
result = spec.get_dummy_input()
self.assertIsInstance(result, set,
"Spec should return a set of tensors")
self.assertEqual(2, len(result), "Spec should return two tensors")
def test_map_input(self):
inp = {'inp': torch.ones(size=(32, 3, 32, 32), dtype=torch.float64)}
model = torch.nn.Sequential(torch.nn.Linear(28 * 28, 100))
spec = FeInputSpec(inp, model)
result = spec.get_dummy_input()
self.assertIsInstance(result, dict,
"Spec should return a tuple of tensors")
self.assertEqual(1, len(result), "Spec should return two tensors")
self.assertIsInstance(result['inp'], torch.Tensor,
"Spec should return a torch tensor")
self.assertListEqual([32, 3, 32, 32], list(result['inp'].shape),
"Result shape is incorrect")
self.assertEqual(torch.float64, result['inp'].dtype,
"Result dtype is incorrect")
def _build_estimator(model: Union[tf.keras.Model, torch.nn.Module], trace: Traceability, axis: int = -1):
train_data, eval_data = mnist.load_data()
test_data = eval_data.split(0.5)
batch_size = 32
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", axis=axis), Minmax(inputs="x", outputs="x")])
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
LRScheduler(model=model, lr_fn=lambda step: cosine_decay(step, cycle_length=3750, init_lr=1e-3)),
trace
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=1,
traces=traces,
max_train_steps_per_epoch=1,
max_eval_steps_per_epoch=None)
fake_data = tf.ones(shape=(batch_size, 28, 28, 1)) if axis == -1 else torch.ones(size=(batch_size, 1, 28, 28))
model.fe_input_spec = FeInputSpec(fake_data, model)
return estimator
def forward(self, data: Union[Tensor, List[Tensor]], state: Dict[str, Any]) -> Union[Tensor, List[Tensor]]:
training = state['mode'] == "train" and self.trainable
if self.epoch_spec != state['epoch']:
# Gather model input specs for the sake of TensorBoard and Traceability
self.model.fe_input_spec = FeInputSpec(data, self.model)
self.epoch_spec = state['epoch']
data = feed_forward(self.model, data, training=training)
return data
def test_list_input(self):
inp = [
torch.ones(size=(32, 3, 32, 32), dtype=torch.float64),
torch.ones(size=(32, 10), dtype=torch.int8)
]
model = torch.nn.Sequential(torch.nn.Linear(28 * 28, 100))
spec = FeInputSpec(inp, model)
result = spec.get_dummy_input()
self.assertIsInstance(result, list,
"Spec should return a list of tensors")
self.assertEqual(2, len(result), "Spec should return two tensors")
self.assertIsInstance(result[0], torch.Tensor,
"Spec should return a torch tensor")
self.assertListEqual([32, 3, 32, 32], list(result[0].shape),
"Result shape is incorrect")
self.assertEqual(torch.float64, result[0].dtype,
"Result dtype is incorrect")
self.assertIsInstance(result[1], torch.Tensor,
"Spec should return a torch tensor")
self.assertListEqual([32, 10], list(result[1].shape),
"Result shape is incorrect")
self.assertEqual(torch.int8, result[1].dtype,
"Result dtype is incorrect")
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 forward(self, data: Union[Tensor, List[Tensor]],
state: Dict[str, Any]) -> Union[Tensor, List[Tensor]]:
training = state['mode'] == "train" and self.trainable
if isinstance(self.model,
torch.nn.Module) and self.epoch_spec != state['epoch']:
# Gather model input specs for the sake of TensorBoard and Traceability
self.model.fe_input_spec = FeInputSpec(data, self.model)
self.epoch_spec = state['epoch']
if self.multi_inputs:
data = feed_forward(self.model, *data, training=training)
else:
data = feed_forward(self.model, data, training=training)
intermediate_outputs = []
for output in self.intermediate_outputs:
intermediate_outputs.append(_unpack_output(output, self.device))
output.clear(
) # This will only help with pytorch memory, tf tensors will remain until next forward
if intermediate_outputs:
data = to_list(data) + intermediate_outputs
return data