def pretrain_model(epochs, batch_size, max_train_steps_per_epoch, save_dir):
# step 1: prepare dataset
train_data, test_data = load_data()
pipeline = fe.Pipeline(
train_data=train_data,
batch_size=batch_size,
ops=[
PadIfNeeded(min_height=40, min_width=40, image_in="x", image_out="x"),
# augmentation 1
RandomCrop(32, 32, image_in="x", image_out="x_aug"),
Sometimes(HorizontalFlip(image_in="x_aug", image_out="x_aug"), prob=0.5),
Sometimes(
ColorJitter(inputs="x_aug", outputs="x_aug", brightness=0.8, contrast=0.8, saturation=0.8, hue=0.2),
prob=0.8),
Sometimes(ToGray(inputs="x_aug", outputs="x_aug"), prob=0.2),
Sometimes(GaussianBlur(inputs="x_aug", outputs="x_aug", blur_limit=(3, 3), sigma_limit=(0.1, 2.0)),
prob=0.5),
ToFloat(inputs="x_aug", outputs="x_aug"),
# augmentation 2
RandomCrop(32, 32, image_in="x", image_out="x_aug2"),
Sometimes(HorizontalFlip(image_in="x_aug2", image_out="x_aug2"), prob=0.5),
Sometimes(
ColorJitter(inputs="x_aug2", outputs="x_aug2", brightness=0.8, contrast=0.8, saturation=0.8, hue=0.2),
prob=0.8),
Sometimes(ToGray(inputs="x_aug2", outputs="x_aug2"), prob=0.2),
Sometimes(GaussianBlur(inputs="x_aug2", outputs="x_aug2", blur_limit=(3, 3), sigma_limit=(0.1, 2.0)),
prob=0.5),
ToFloat(inputs="x_aug2", outputs="x_aug2")
])
# step 2: prepare network
model_con, model_finetune = fe.build(model_fn=ResNet9, optimizer_fn=["adam", "adam"])
network = fe.Network(ops=[
LambdaOp(lambda x, y: tf.concat([x, y], axis=0), inputs=["x_aug", "x_aug2"], outputs="x_com"),
ModelOp(model=model_con, inputs="x_com", outputs="y_com"),
LambdaOp(lambda x: tf.split(x, 2, axis=0), inputs="y_com", outputs=["y_pred", "y_pred2"]),
NTXentOp(arg1="y_pred", arg2="y_pred2", outputs=["NTXent", "logit", "label"]),
UpdateOp(model=model_con, loss_name="NTXent")
])
# step 3: prepare estimator
traces = [
Accuracy(true_key="label", pred_key="logit", mode="train", output_name="contrastive_accuracy"),
ModelSaver(model=model_con, save_dir=save_dir),
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
monitor_names="contrastive_accuracy")
estimator.fit()
return model_con, model_finetune
def test_single_input_tf_static(self):
a = LambdaOp(inputs='x', outputs='x', fn=lambda x: x + 1)
b = LambdaOp(inputs='x', outputs='x', fn=lambda x: x + 5)
oneof = OneOf(a, b)
oneof.build('tf')
output = oneof.forward(data=self.single_input_tf, state={})
with self.subTest('Check output type'):
self.assertTrue(tf.is_tensor(output))
with self.subTest('Check output image shape'):
self.assertEqual(output.shape, self.output_shape)
def test_single_input_torch(self):
a = LambdaOp(inputs='x', outputs='x', fn=lambda x: x + 1)
b = LambdaOp(inputs='x', outputs='x', fn=lambda x: x + 5)
oneof = OneOf(a, b)
oneof.build('torch')
output = oneof.forward(data=self.single_input_torch, state={})
with self.subTest('Check output type'):
self.assertEqual(type(output), torch.Tensor)
with self.subTest('Check output image shape'):
self.assertEqual(output.shape, self.output_shape)
def test_single_input_tf_static(self):
a = LambdaOp(inputs='x', outputs='y', fn=lambda x: x + 1, mode='test')
b = LambdaOp(inputs=['y', 'z'], outputs='w', fn=lambda x, y: x + y, mode='test')
fuse = Fuse([a, b])
with self.subTest('Check op inputs'):
self.assertListEqual(fuse.inputs, ['x', 'z'])
with self.subTest('Check op outputs'):
self.assertListEqual(fuse.outputs, ['y', 'w'])
with self.subTest('Check op mode'):
self.assertSetEqual(fuse.mode, {'test'})
output = fuse.forward(data=self.multi_input_tf, state={"mode": "test", "deferred": {}})
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output image shape'):
self.assertEqual(output[0].shape, self.output_shape)
def test_repeat_fn_exterior_value_static(self):
add_op = LambdaOp(inputs='x',
outputs=('x', 'y'),
fn=lambda w: (w + 1, w * w),
mode='eval')
repeat_op = Repeat(add_op, repeat=lambda y, z: y + z < 25)
repeat_op.build('tf')
with self.subTest('Check op inputs'):
self.assertListEqual(repeat_op.inputs, ['x', 'z'])
with self.subTest('Check op outputs'):
self.assertListEqual(repeat_op.outputs, ['x', 'y'])
with self.subTest('Check op mode'):
self.assertSetEqual(repeat_op.mode, {'eval'})
x = [tf.ones([1]), 10 + tf.ones([1])]
output = tf.function(lambda y: repeat_op.forward(data=y,
state={
"deferred": {},
"mode": "eval"
}))
output(x)
output = output(x)
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output value (x)'):
self.assertEqual(5, output[0])
with self.subTest('Check output value (y)'):
self.assertEqual(16, output[1])
def test_single_repeat_fn_interior_value_static(self):
add_op = LambdaOp(inputs='x',
outputs=('x', 'y'),
fn=lambda z: (z + 1, z * z),
mode='eval')
repeat_op = Repeat(add_op, repeat=lambda y: y < 1)
repeat_op.build('tf')
with self.subTest('Check op inputs'):
self.assertListEqual(repeat_op.inputs, ['x'])
with self.subTest('Check op outputs'):
self.assertListEqual(repeat_op.outputs, ['x', 'y'])
with self.subTest('Check op mode'):
self.assertSetEqual(repeat_op.mode, {'eval'})
x = [tf.ones([1])]
output = tf.function(lambda y: repeat_op.forward(data=y,
state={
"deferred": {},
"mode": "eval"
}))
output(x) # build the graph
output = output(x)
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output value (x)'):
self.assertEqual(2, output[0])
with self.subTest('Check output value (y)'):
self.assertEqual(1, output[1])
def test_repeat_fn_exterior_value_tf(self):
add_op = LambdaOp(inputs='x',
outputs=('x', 'y'),
fn=lambda x: (x + 1, x * x),
mode='eval')
repeat_op = Repeat(add_op, repeat=lambda y, z: y + z < 25)
repeat_op.build('tf')
with self.subTest('Check op inputs'):
self.assertListEqual(repeat_op.inputs, ['x', 'z'])
with self.subTest('Check op outputs'):
self.assertListEqual(repeat_op.outputs, ['x', 'y'])
with self.subTest('Check op mode'):
self.assertSetEqual(repeat_op.mode, {'eval'})
with tf.GradientTape(persistent=True) as tape:
output = repeat_op.forward(data=[tf.ones([1]), 10 + tf.ones([1])],
state={
"deferred": {},
"mode": "eval",
"tape": tape
})
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output value (x)'):
self.assertEqual(5, output[0])
with self.subTest('Check output value (y)'):
self.assertEqual(16, output[1])
def get_estimator(epochs=2, batch_size=32):
# step 1
train_data, eval_data = mnist.load_data()
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
batch_size=batch_size,
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", "feature_vector"],
intermediate_layers='dense'),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
CustomLoss(inputs=("feature_vector", "feature_selected"),
outputs="feature_loss"),
LambdaOp(fn=lambda x, y: x + y,
inputs=("ce", "feature_loss"),
outputs="total_loss"),
UpdateOp(model=model, loss_name="total_loss")
])
# step 3
traces = [
MemoryBank(inputs=("feature_vector", "y"), outputs="feature_selected")
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces)
return estimator
def test_single_input_torch(self):
a = LambdaOp(inputs='x', outputs='x', fn=lambda x: x + 1)
sometimes = Sometimes(a, prob=0.75)
sometimes.build('torch')
output = sometimes.forward(data=self.single_input_torch, state={})
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output image shape'):
self.assertEqual(output[0].shape, self.output_shape)
def test_multi_input_tf_static(self):
a = LambdaOp(inputs=['x', 'y'],
outputs=['y', 'z'],
fn=lambda x, y: [x + y, x - y])
b = LambdaOp(inputs=['x', 'y'],
outputs=['y', 'z'],
fn=lambda x, y: [x * y, x + y])
c = LambdaOp(inputs=['x', 'y'],
outputs=['y', 'z'],
fn=lambda x, y: [y, x])
oneof = OneOf(a, b, c)
oneof.build('tf')
output = oneof.forward(data=self.multi_input_tf, state={})
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output list length'):
self.assertEqual(len(output), 2)
for img_output in output:
with self.subTest('Check output image shape'):
self.assertEqual(img_output.shape, self.output_shape)
def test_multi_input_tf_static(self):
a = LambdaOp(inputs=['x', 'y'],
outputs=['y', 'x'],
fn=lambda x, y: [x + y, x - y])
sometimes = Sometimes(a)
sometimes.build('tf')
output = sometimes.forward(data=self.multi_input_tf, state={})
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output list length'):
self.assertEqual(len(output), 2)
for img_output in output:
with self.subTest('Check output image shape'):
self.assertEqual(img_output.shape, self.output_shape)
def test_multi_repeat_fn_interior_value_tf(self):
add_op = LambdaOp(inputs='x', outputs=('x', 'y'), fn=lambda x: (x + 1, x * x), mode='eval')
repeat_op = Repeat(add_op, repeat=lambda y: y < 25)
repeat_op.build('tf')
with self.subTest('Check op inputs'):
self.assertListEqual(repeat_op.inputs, ['x'])
with self.subTest('Check op outputs'):
self.assertListEqual(repeat_op.outputs, ['x', 'y'])
with self.subTest('Check op mode'):
self.assertSetEqual(repeat_op.mode, {'eval'})
output = repeat_op.forward(data=[tf.ones([1])], state={"deferred": {}, "mode": "eval"})
with self.subTest('Check output type'):
self.assertEqual(type(output), list)
with self.subTest('Check output value (x)'):
self.assertEqual(6, output[0])
with self.subTest('Check output value (y)'):
self.assertEqual(25, output[1])
def pretrain_model(epochs, batch_size, train_steps_per_epoch, save_dir):
train_data, test_data = load_data()
pipeline = fe.Pipeline(
train_data=train_data,
batch_size=batch_size,
ops=[
PadIfNeeded(min_height=40,
min_width=40,
image_in="x",
image_out="x",
mode="train"),
# augmentation 1
RandomCrop(32, 32, image_in="x", image_out="x_aug"),
Sometimes(HorizontalFlip(image_in="x_aug", image_out="x_aug"),
prob=0.5),
Sometimes(ColorJitter(inputs="x_aug",
outputs="x_aug",
brightness=0.8,
contrast=0.8,
saturation=0.8,
hue=0.2),
prob=0.8),
Sometimes(ToGray(inputs="x_aug", outputs="x_aug"), prob=0.2),
Sometimes(GaussianBlur(inputs="x_aug",
outputs="x_aug",
blur_limit=(3, 3),
sigma_limit=(0.1, 2.0)),
prob=0.5),
ChannelTranspose(inputs="x_aug", outputs="x_aug"),
ToFloat(inputs="x_aug", outputs="x_aug"),
# augmentation 2
RandomCrop(32, 32, image_in="x", image_out="x_aug2"),
Sometimes(HorizontalFlip(image_in="x_aug2", image_out="x_aug2"),
prob=0.5),
Sometimes(ColorJitter(inputs="x_aug2",
outputs="x_aug2",
brightness=0.8,
contrast=0.8,
saturation=0.8,
hue=0.2),
prob=0.8),
Sometimes(ToGray(inputs="x_aug2", outputs="x_aug2"), prob=0.2),
Sometimes(GaussianBlur(inputs="x_aug2",
outputs="x_aug2",
blur_limit=(3, 3),
sigma_limit=(0.1, 2.0)),
prob=0.5),
ChannelTranspose(inputs="x_aug2", outputs="x_aug2"),
ToFloat(inputs="x_aug2", outputs="x_aug2")
])
model_con = fe.build(model_fn=lambda: ResNet9OneLayerHead(length=128),
optimizer_fn="adam")
network = fe.Network(ops=[
LambdaOp(lambda x, y: torch.cat([x, y], dim=0),
inputs=["x_aug", "x_aug2"],
outputs="x_com"),
ModelOp(model=model_con, inputs="x_com", outputs="y_com"),
LambdaOp(lambda x: torch.chunk(x, 2, dim=0),
inputs="y_com",
outputs=["y_pred", "y_pred2"],
mode="train"),
NTXentOp(arg1="y_pred",
arg2="y_pred2",
outputs=["NTXent", "logit", "label"],
mode="train"),
UpdateOp(model=model_con, loss_name="NTXent")
])
traces = [
Accuracy(true_key="label",
pred_key="logit",
mode="train",
output_name="contrastive_accuracy"),
ModelSaver(model=model_con, save_dir=save_dir)
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch)
estimator.fit()
return model_con
def test_multi_input(self):
op = LambdaOp(fn=tf.reshape)
data = op.forward(data=[tf.convert_to_tensor([1, 2, 3, 4]), (2, 2)],
state={})
self.assertTrue(is_equal(data, tf.convert_to_tensor([[1, 2], [3, 4]])))
def test_single_input(self):
op = LambdaOp(fn=tf.reduce_sum)
data = op.forward(data=tf.convert_to_tensor([[1, 2, 3]]), state={})
self.assertEqual(data, 6)