def test_tf_input(self):
gradient = GradientOp(inputs='x', finals='x', outputs='y')
with tf.GradientTape(persistent=True) as tape:
x = self.tf_data * self.tf_data
output = gradient.forward(data=[self.tf_data, x],
state={'tape': tape})
self.assertTrue(is_equal(output, self.tf_output))
def test_torch_input(self):
gradient = GradientOp(inputs='x', finals='x', outputs='y')
gradient.build("torch")
x = self.torch_data * self.torch_data
output = gradient.forward(data=[self.torch_data, x],
state={'tape': None})
self.assertTrue(is_equal(output, self.torch_output))
def test_torch_model_input(self):
gradient = GradientOp(finals="pred",
outputs="grad",
model=self.torch_model)
with tf.GradientTape(persistent=True) as tape:
pred = feed_forward(model=self.torch_model,
x=torch.tensor([[1.0, 1.0, 1.0],
[1.0, -1.0, -0.5]]))
output = gradient.forward(data=[pred], state={"tape": tape})
self.assertTrue(
is_equal(output[0][0].numpy(),
np.array([[2.0, 0.0, 0.5]], dtype="float32")))
def test_tf_input(self):
def update_gradient(gradient):
with tf.GradientTape(persistent=True) as tape:
x = self.tf_data * self.tf_data
output = gradient.forward(data=[self.tf_data, x],
state={'tape': tape})
self.assertTrue(is_equal(output, self.tf_output))
gradient = GradientOp(inputs='x', finals='x', outputs='y')
gradient.build("tf")
strategy = tf.distribute.get_strategy()
if isinstance(strategy, tf.distribute.MirroredStrategy):
strategy.run(update_gradient, args=(gradient, ))
else:
update_gradient(gradient)
def run_test(mixed_precision, merge_grad, gradient):
lr = 0.1
lr2 = 0.01
pipeline = fe.Pipeline(train_data=self.train_data,
batch_size=4,
ops=[ExpandDims(inputs="x", outputs="x", axis=0), Minmax(inputs="x", outputs="x")])
optimizer_fn = RepeatScheduler(
[lambda x: torch.optim.SGD(params=x, lr=lr), lambda x: torch.optim.SGD(params=x, lr=lr2)])
model = fe.build(model_fn=LeNet_torch, optimizer_fn=optimizer_fn, mixed_precision=mixed_precision)
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
GradientOp(model=model, finals="ce", outputs="grad"),
UpdateOp(model=model, loss_name="ce", gradients=gradient, merge_grad=merge_grad),
])
traces = [
CheckNetworkWeight(model=model,
grad_key="grad",
merge_grad=merge_grad,
test_self=self,
framework="torch",
lrs=[lr, lr2, lr, lr2],
work_intervals=[[1, 2], [2, 3], [3, 4], [4, 5]])
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=4,
traces=traces,
train_steps_per_epoch=2)
estimator.fit(warmup=False)
def run_test(mixed_precision, merge_grad, gradient):
lr = 0.1
pipeline = fe.Pipeline(train_data=self.train_data,
batch_size=4,
ops=[ExpandDims(inputs="x", outputs="x"), Minmax(inputs="x", outputs="x")])
model = fe.build(model_fn=LeNet_tf,
optimizer_fn=lambda: tf.optimizers.SGD(lr),
mixed_precision=mixed_precision)
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
GradientOp(model=model, finals="ce", outputs="grad"),
UpdateOp(model=model, loss_name="ce", gradients=gradient, merge_grad=merge_grad),
])
traces = [
CheckNetworkWeight(model=model,
grad_key="grad",
merge_grad=merge_grad,
test_self=self,
lrs=lr,
framework="tf")
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=2,
traces=traces,
train_steps_per_epoch=2)
estimator.fit(warmup=False)
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_model_input(self):
def update_gradient(gradient):
with tf.GradientTape(persistent=True) as tape:
pred = feed_forward(
self.tf_model,
tf.constant([[1.0, 1.0, 1.0], [1.0, -1.0, -0.5]]))
output = gradient.forward(data=[pred], state={"tape": tape})
self.assertTrue(
is_equal(output[0][0].numpy(),
np.array([[2.0], [0.0], [0.5]], dtype="float32")))
gradient = GradientOp(finals="pred",
outputs="grad",
model=self.tf_model)
gradient.build("tf")
strategy = tf.distribute.get_strategy()
if isinstance(strategy, tf.distribute.MirroredStrategy):
strategy.run(update_gradient, args=(gradient, ))
else:
update_gradient(gradient)
def test_tf_model_end_to_end_gradient(self):
train_data, _ = mnist.load_data()
pipeline = fe.Pipeline(train_data=train_data,
batch_size=4,
ops=[ExpandDims(inputs="x", outputs="x"), Minmax(inputs="x", outputs="x")])
model = fe.build(model_fn=LeNet_tf, optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
GradientOp(model=model, finals="ce", outputs="gradients"),
UpdateOp(model=model, gradients="gradients", loss_name="ce")
])
estimator = fe.Estimator(pipeline=pipeline, network=network, epochs=2, max_train_steps_per_epoch=10)
estimator.fit()
def get_estimator():
# step 1
train_data, eval_data = mnist.load_data()
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
batch_size=32,
ops=[
ExpandDims(inputs="x", outputs="x"),
Minmax(inputs="x", outputs="x")
])
# step 2
model = fe.build(model_fn=LeNet, optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
GradientOp(finals="ce", model=model, outputs="grad"),
ClipGrad(inputs="grad", outputs="grad", mode="train"),
UpdateOp(model=model, gradients="grad", loss_name="ce")
])
# step 3
estimator = fe.Estimator(pipeline=pipeline, network=network, epochs=2)
return estimator