def test_single_input(self):
channel_transpose = ChannelTranspose(inputs='x', outputs='x')
output = channel_transpose.forward(data=self.single_input, 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.single_output_shape)
def test_multi_input(self):
channel_transpose = ChannelTranspose(inputs='x', outputs='x')
output = channel_transpose.forward(data=self.multi_input, 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 mask shape'):
self.assertEqual(img_output.shape, self.multi_output_shape)
def finetune_model(model, epochs, batch_size, train_steps_per_epoch, save_dir):
train_data, test_data = load_data()
train_data = train_data.split(0.1)
pipeline = fe.Pipeline(train_data=train_data,
eval_data=test_data,
batch_size=batch_size,
ops=[
ToFloat(inputs="x", outputs="x"),
ChannelTranspose(inputs="x", outputs="x"),
])
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=["y_pred", "y"], outputs="ce", from_logits=True),
UpdateOp(model=model, loss_name="ce")
])
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch)
estimator.fit()
def get_estimator(epochs=50,
batch_size=128,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1
train_data, eval_data = cifair100.load_data()
# Add label noise to simulate real-world labeling problems
corrupt_dataset(train_data)
test_data = eval_data.split(range(len(eval_data) // 2))
pipeline = fe.Pipeline(
train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x", outputs="x", mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616)),
PadIfNeeded(min_height=40, min_width=40, image_in="x", image_out="x", mode="train"),
RandomCrop(32, 32, image_in="x", image_out="x", mode="train"),
Sometimes(HorizontalFlip(image_in="x", image_out="x", mode="train")),
CoarseDropout(inputs="x", outputs="x", mode="train", max_holes=1),
ChannelTranspose(inputs="x", outputs="x")
])
# step 2
model = fe.build(model_fn=big_lenet, optimizer_fn='adam')
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
SuperLoss(CrossEntropy(inputs=("y_pred", "y"), outputs="ce"), output_confidence="confidence"),
UpdateOp(model=model, loss_name="ce")
])
# step 3
traces = [
MCC(true_key="y", pred_key="y_pred"),
BestModelSaver(model=model, save_dir=save_dir, metric="mcc", save_best_mode="max", load_best_final=True),
LabelTracker(metric="confidence",
label="data_labels",
label_mapping={
"Normal": 0, "Corrupted": 1
},
mode="train",
outputs="label_confidence"),
ImageSaver(inputs="label_confidence", save_dir=save_dir, mode="train"),
]
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 get_estimator(epsilon=0.04,
epochs=10,
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 = cifair10.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=[
Normalize(inputs="x",
outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
ChannelTranspose(inputs="x", outputs="x"),
])
# step 2
model = fe.build(model_fn=lambda: LeNet(input_shape=(3, 32, 32)),
optimizer_fn="adam")
network = fe.Network(ops=[
Watch(inputs="x"),
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="base_ce"),
FGSM(data="x", loss="base_ce", outputs="x_adverse", epsilon=epsilon),
ModelOp(model=model, inputs="x_adverse", outputs="y_pred_adv"),
CrossEntropy(inputs=("y_pred_adv", "y"), outputs="adv_ce"),
Average(inputs=("base_ce", "adv_ce"), outputs="avg_ce"),
UpdateOp(model=model, loss_name="avg_ce")
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name="base accuracy"),
Accuracy(true_key="y",
pred_key="y_pred_adv",
output_name="adversarial accuracy"),
BestModelSaver(model=model,
save_dir=save_dir,
metric="adv_ce",
save_best_mode="min"),
]
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,
monitor_names=["base_ce", "adv_ce"])
return estimator
def get_estimator(batch_size=4,
epochs=2,
max_train_steps_per_epoch=None,
log_steps=100,
style_weight=5.0,
content_weight=1.0,
tv_weight=1e-4,
save_dir=tempfile.mkdtemp(),
style_img_path='Vassily_Kandinsky,_1913_-_Composition_7.jpg',
data_dir=None):
train_data, _ = mscoco.load_data(root_dir=data_dir, load_bboxes=False, load_masks=False, load_captions=False)
device = "cuda" if torch.cuda.is_available() else "cpu"
style_img = cv2.imread(style_img_path)
assert style_img is not None, "cannot load the style image, please go to the folder with style image"
style_img = cv2.resize(style_img, (256, 256))
style_img = (style_img.astype(np.float32) - 127.5) / 127.5
pipeline = fe.Pipeline(
train_data=train_data,
batch_size=batch_size,
ops=[
ReadImage(inputs="image", outputs="image"),
Normalize(inputs="image", outputs="image", mean=1.0, std=1.0, max_pixel_value=127.5),
Resize(height=256, width=256, image_in="image", image_out="image"),
LambdaOp(fn=lambda: style_img, outputs="style_image"),
ChannelTranspose(inputs=["image", "style_image"], outputs=["image", "style_image"])
])
model = fe.build(model_fn=StyleTransferNet,
model_name="style_transfer_net",
optimizer_fn=lambda x: torch.optim.Adam(x, lr=1e-3))
network = fe.Network(ops=[
ModelOp(inputs="image", model=model, outputs="image_out"),
ExtractVGGFeatures(inputs="style_image", outputs="y_style", device=device),
ExtractVGGFeatures(inputs="image", outputs="y_content", device=device),
ExtractVGGFeatures(inputs="image_out", outputs="y_pred", device=device),
StyleContentLoss(style_weight=style_weight,
content_weight=content_weight,
tv_weight=tv_weight,
inputs=('y_pred', 'y_style', 'y_content', 'image_out'),
outputs='loss'),
UpdateOp(model=model, loss_name="loss")
])
estimator = fe.Estimator(network=network,
pipeline=pipeline,
traces=ModelSaver(model=model, save_dir=save_dir, frequency=1),
epochs=epochs,
max_train_steps_per_epoch=max_train_steps_per_epoch,
log_steps=log_steps)
return estimator
def get_estimator(epochs=24, batch_size=128, lr_epochs=100, max_train_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1: prepare dataset
train_data, test_data = load_data()
pipeline = fe.Pipeline(
train_data=train_data,
eval_data=test_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x", outputs="x", mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616)),
PadIfNeeded(min_height=40, min_width=40, image_in="x", image_out="x", mode="train"),
RandomCrop(32, 32, image_in="x", image_out="x", mode="train"),
Sometimes(HorizontalFlip(image_in="x", image_out="x", mode="train")),
CoarseDropout(inputs="x", outputs="x", mode="train", max_holes=1),
ChannelTranspose(inputs="x", outputs="x"),
Onehot(inputs="y", outputs="y", mode="train", num_classes=10, label_smoothing=0.2)
])
# step 2: prepare network
model = fe.build(model_fn=ResNet9, optimizer_fn="sgd")
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")
])
# get the max learning rate
lr_max = search_max_lr(pipeline=pipeline, model=model, network=network, epochs=lr_epochs)
lr_min = lr_max / 40
print(f"The maximum LR: {lr_max}, and minimun LR: {lr_min}")
mid_step = int(epochs * 0.45 * len(train_data) / batch_size)
end_step = int(epochs * len(train_data) / batch_size)
# reinitialize the model
model = fe.build(model_fn=ResNet9, optimizer_fn="sgd")
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")
])
# step 3: prepare estimator
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
BestModelSaver(model=model, save_dir=save_dir, metric="accuracy", save_best_mode="max"),
LRScheduler(model=model, lr_fn=lambda step: super_schedule(step, lr_max, lr_min, mid_step, end_step))
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch)
return estimator
def get_estimator(epsilon=0.04,
epochs=20,
batch_size=32,
code_length=16,
train_steps_per_epoch=None,
eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1
train_data, eval_data = cifair10.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=[
Normalize(inputs="x", outputs="x", mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616)),
ChannelTranspose(inputs="x", outputs="x"),
Hadamard(inputs="y", outputs="y_code", n_classes=10)
],
num_process=0)
# step 2
model = fe.build(model_fn=lambda: EccLeNet(code_length=code_length), optimizer_fn="adam")
network = fe.Network(ops=[
Watch(inputs="x", mode=('eval', 'test')),
ModelOp(model=model, inputs="x", outputs="y_pred_code"),
Hinge(inputs=("y_pred_code", "y_code"), outputs="base_hinge"),
UpdateOp(model=model, loss_name="base_hinge"),
UnHadamard(inputs="y_pred_code", outputs="y_pred", n_classes=10, mode=('eval', 'test')),
# The adversarial attack:
FGSM(data="x", loss="base_hinge", outputs="x_adverse_hinge", epsilon=epsilon, mode=('eval', 'test')),
ModelOp(model=model, inputs="x_adverse_hinge", outputs="y_pred_adv_hinge_code", mode=('eval', 'test')),
Hinge(inputs=("y_pred_adv_hinge_code", "y_code"), outputs="adv_hinge", mode=('eval', 'test')),
UnHadamard(inputs="y_pred_adv_hinge_code", outputs="y_pred_adv_hinge", n_classes=10, mode=('eval', 'test')),
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name="base_accuracy"),
Accuracy(true_key="y", pred_key="y_pred_adv_hinge", output_name="adversarial_accuracy"),
BestModelSaver(model=model, save_dir=save_dir, metric="base_hinge", save_best_mode="min", load_best_final=True)
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch,
monitor_names=["adv_hinge"])
return estimator
def get_pipeline_data(batch_size=128):
train_data, _ = cifar10.load_data()
pipeline = fe.Pipeline(train_data=train_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x",
outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
ChannelTranspose(inputs="x", outputs="x")
])
result = pipeline.get_results()
return result
def get_estimator(level,
num_augment,
epochs=24,
batch_size=512,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None):
assert 0 <= level <= 10, "the level should be between 0 and 10"
train_data, test_data = load_data()
aug_ops = [
OneOf(
Rotate(level=level, inputs="x", outputs="x", mode="train"),
Identity(level=level, inputs="x", outputs="x", mode="train"),
AutoContrast(level=level, inputs="x", outputs="x", mode="train"),
Equalize(level=level, inputs="x", outputs="x", mode="train"),
Posterize(level=level, inputs="x", outputs="x", mode="train"),
Solarize(level=level, inputs="x", outputs="x", mode="train"),
Sharpness(level=level, inputs="x", outputs="x", mode="train"),
Contrast(level=level, inputs="x", outputs="x", mode="train"),
Color(level=level, inputs="x", outputs="x", mode="train"),
Brightness(level=level, inputs="x", outputs="x", mode="train"),
ShearX(level=level, inputs="x", outputs="x", mode="train"),
ShearY(level=level, inputs="x", outputs="x", mode="train"),
TranslateX(level=level, inputs="x", outputs="x", mode="train"),
TranslateY(level=level, inputs="x", outputs="x", mode="train"),
) for _ in range(num_augment)
]
pipeline = fe.Pipeline(train_data=train_data,
test_data=test_data,
batch_size=batch_size,
ops=aug_ops + [
Normalize(inputs="x",
outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
ChannelTranspose(inputs="x", outputs="x"),
])
model = fe.build(model_fn=MyModel, optimizer_fn="adam")
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")
])
estimator = fe.Estimator(
pipeline=pipeline,
network=network,
epochs=epochs,
traces=Accuracy(true_key="y", pred_key="y_pred"),
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch)
return estimator
def __init__(self, data_path, batch_size, collaborator_count, **kwargs):
"""
Initialize.
Args:
data_path: File path for the dataset
batch_size (int): The batch size for the data loader
**kwargs: Additional arguments, passed to super init and load_mnist_shard
"""
# TODO: We should be downloading the dataset shard into a directory
# TODO: There needs to be a method to ask how many collaborators and
# what index/rank is this collaborator.
# Then we have a way to automatically shard based on rank and size of
# collaborator list.
train_data, eval_data = cifar10.load_data()
test_data = eval_data.split(0.5)
train_data, eval_data, test_data = self.split_data(
train_data,
eval_data,
test_data,
int(data_path),
collaborator_count
)
super().__init__(fe.Pipeline(
train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x", outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
ChannelTranspose(inputs="x", outputs="x")
]), **kwargs)
print(f"train_data = {train_data}")
print(f"eval_data = {eval_data}")
print(f"test_data = {test_data}")
print(f"batch_size = {batch_size}")
def get_estimator(epochs=24, batch_size=512, max_train_steps_per_epoch=None, save_dir=tempfile.mkdtemp()):
# step 1: prepare dataset
train_data, test_data = load_data()
pipeline = fe.Pipeline(
train_data=train_data,
test_data=test_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x", outputs="x", mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616)),
PadIfNeeded(min_height=40, min_width=40, image_in="x", image_out="x", mode="train"),
RandomCrop(32, 32, image_in="x", image_out="x", mode="train"),
Sometimes(HorizontalFlip(image_in="x", image_out="x", mode="train")),
CoarseDropout(inputs="x", outputs="x", mode="train", max_holes=1),
ChannelTranspose(inputs="x", outputs="x"),
Onehot(inputs="y", outputs="y", mode="train", num_classes=10, label_smoothing=0.2)
])
# step 2: prepare network
model = fe.build(model_fn=FastCifar, optimizer_fn="adam")
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")
])
# step 3 prepare estimator
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
BestModelSaver(model=model, save_dir=save_dir, metric="accuracy", save_best_mode="max"),
LRScheduler(model=model, lr_fn=lr_schedule)
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch)
return estimator
def get_estimator(epochs=150,
batch_size=32,
save_dir=tempfile.mkdtemp(),
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None):
# step 1: prepare dataset
train_data, eval_data = load_data()
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
batch_size=batch_size * get_num_devices(),
ops=[
Normalize(inputs="x",
outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
PadIfNeeded(min_height=40,
min_width=40,
image_in="x",
image_out="x",
mode="train"),
RandomCrop(32,
32,
image_in="x",
image_out="x",
mode="train"),
Sometimes(
HorizontalFlip(image_in="x",
image_out="x",
mode="train")),
CoarseDropout(inputs="x",
outputs="x",
mode="train",
max_holes=1),
ChannelTranspose(inputs="x", outputs="x"),
])
# step 2: prepare network
model = fe.build(model_fn=lambda: PyramidNet(
depth=272, alpha=200, num_classes=10, bottleneck=True),
optimizer_fn=lambda x: torch.optim.SGD(
x, lr=0.1, momentum=0.9, weight_decay=0.0001))
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce", from_logits=True),
UpdateOp(model=model, loss_name="ce")
])
# step 3 prepare estimator
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
LRScheduler(model=model, lr_fn=lr_schedule),
BestModelSaver(model=model,
save_dir=save_dir,
metric="accuracy",
save_best_mode="max")
]
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 get_estimator(epochs=12, batch_size=512, save_dir=tempfile.mkdtemp()):
# epoch 1-10, train on cifair100, epoch 11-end: train on cifar10
cifair10_train, cifari10_test = cifair10.load_data()
cifair100_train, _ = cifair100.load_data()
train_ds = EpochScheduler({1: cifair100_train, 11: cifair10_train})
pipeline = fe.Pipeline(train_data=train_ds,
test_data=cifari10_test,
batch_size=batch_size,
ops=[
Normalize(inputs="x",
outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
PadIfNeeded(min_height=40,
min_width=40,
image_in="x",
image_out="x",
mode="train"),
RandomCrop(32,
32,
image_in="x",
image_out="x",
mode="train"),
Sometimes(
HorizontalFlip(image_in="x",
image_out="x",
mode="train")),
CoarseDropout(inputs="x",
outputs="x",
mode="train",
max_holes=1),
ChannelTranspose(inputs="x", outputs="x")
])
# step 2: prepare network
backbone = fe.build(model_fn=lambda: Backbone(input_size=(3, 32, 32)),
optimizer_fn="adam")
cls_head_cifar100 = fe.build(model_fn=lambda: Classifier(classes=100),
optimizer_fn="adam")
cls_head_cifar10 = fe.build(model_fn=lambda: Classifier(classes=10),
optimizer_fn="adam")
# if you want to save the final cifar10 model, you can build a model then provide it to ModelSaver
# final_model_cifar10 = fe.build(model_fn=lambda: MyModel(backbone, cls_head_cifar10), optimizer_fn=None)
# epoch 1-10: train backbone and cls_head_cifar100, epoch 11-end: train cls_head_cifar10 only
ModelOp_cls_head = EpochScheduler({
1:
ModelOp(model=cls_head_cifar100, inputs="feature", outputs="y_pred"),
11:
ModelOp(model=cls_head_cifar10, inputs="feature", outputs="y_pred"),
})
UpdateOp_backbone = EpochScheduler({
1:
UpdateOp(model=backbone, loss_name="ce"),
11:
None
})
UpdateOp_cls_head = EpochScheduler({
1:
UpdateOp(model=cls_head_cifar100, loss_name="ce"),
11:
UpdateOp(model=cls_head_cifar10, loss_name="ce")
})
network = fe.Network(ops=[
ModelOp(model=backbone, inputs="x",
outputs="feature"), ModelOp_cls_head,
CrossEntropy(inputs=("y_pred", "y"), outputs="ce", from_logits=True),
UpdateOp_backbone, UpdateOp_cls_head
])
traces = [Accuracy(true_key="y", pred_key="y_pred")]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces)
return estimator
def get_estimator(data_dir=None,
model_dir=tempfile.mkdtemp(),
batch_size=16,
epochs=13,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
image_size=512,
num_classes=90):
# pipeline
train_ds, eval_ds = mscoco.load_data(root_dir=data_dir)
pipeline = fe.Pipeline(
train_data=train_ds,
eval_data=eval_ds,
batch_size=batch_size,
ops=[
ReadImage(inputs="image", outputs="image"),
LongestMaxSize(image_size,
image_in="image",
image_out="image",
bbox_in="bbox",
bbox_out="bbox",
bbox_params=BboxParams("coco", min_area=1.0)),
PadIfNeeded(
image_size,
image_size,
border_mode=cv2.BORDER_CONSTANT,
image_in="image",
image_out="image",
bbox_in="bbox",
bbox_out="bbox",
bbox_params=BboxParams("coco", min_area=1.0),
),
Sometimes(
HorizontalFlip(mode="train",
image_in="image",
image_out="image",
bbox_in="bbox",
bbox_out="bbox",
bbox_params='coco')),
# normalize from uint8 to [-1, 1]
Normalize(inputs="image",
outputs="image",
mean=1.0,
std=1.0,
max_pixel_value=127.5),
ShiftLabel(inputs="bbox", outputs="bbox"),
AnchorBox(inputs="bbox",
outputs="anchorbox",
width=image_size,
height=image_size),
ChannelTranspose(inputs="image", outputs="image")
],
pad_value=0)
# network
model = fe.build(model_fn=lambda: RetinaNet(num_classes=num_classes),
optimizer_fn=lambda x: torch.optim.SGD(
x, lr=2e-4, momentum=0.9, weight_decay=0.0001))
network = fe.Network(ops=[
ModelOp(model=model, inputs="image", outputs=["cls_pred", "loc_pred"]),
RetinaLoss(inputs=["anchorbox", "cls_pred", "loc_pred"],
outputs=["total_loss", "focal_loss", "l1_loss"]),
UpdateOp(model=model, loss_name="total_loss"),
PredictBox(input_shape=(image_size, image_size, 3),
inputs=["cls_pred", "loc_pred"],
outputs="pred",
mode="eval")
])
# estimator
traces = [
LRScheduler(model=model, lr_fn=lr_fn),
BestModelSaver(model=model,
save_dir=model_dir,
metric='mAP',
save_best_mode="max"),
MeanAveragePrecision(num_classes=num_classes,
true_key='bbox',
pred_key='pred',
mode="eval")
]
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,
monitor_names=["l1_loss", "focal_loss"])
return estimator
def get_estimator(weight=10.0,
epochs=200,
batch_size=1,
max_train_steps_per_epoch=None,
save_dir=tempfile.mkdtemp(),
data_dir=None):
train_data, _ = load_data(batch_size=batch_size, root_dir=data_dir)
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = fe.Pipeline(train_data=train_data,
ops=[
ReadImage(inputs=["A", "B"], outputs=["A",
"B"]),
Normalize(inputs=["A", "B"],
outputs=["real_A", "real_B"],
mean=1.0,
std=1.0,
max_pixel_value=127.5),
Resize(height=286,
width=286,
image_in="real_A",
image_out="real_A",
mode="train"),
RandomCrop(height=256,
width=256,
image_in="real_A",
image_out="real_A",
mode="train"),
Resize(height=286,
width=286,
image_in="real_B",
image_out="real_B",
mode="train"),
RandomCrop(height=256,
width=256,
image_in="real_B",
image_out="real_B",
mode="train"),
Sometimes(
HorizontalFlip(image_in="real_A",
image_out="real_A",
mode="train")),
Sometimes(
HorizontalFlip(image_in="real_B",
image_out="real_B",
mode="train")),
ChannelTranspose(inputs=["real_A", "real_B"],
outputs=["real_A", "real_B"])
])
g_AtoB = fe.build(model_fn=Generator,
model_name="g_AtoB",
optimizer_fn=lambda x: torch.optim.Adam(
x, lr=2e-4, betas=(0.5, 0.999)))
g_BtoA = fe.build(model_fn=Generator,
model_name="g_BtoA",
optimizer_fn=lambda x: torch.optim.Adam(
x, lr=2e-4, betas=(0.5, 0.999)))
d_A = fe.build(model_fn=Discriminator,
model_name="d_A",
optimizer_fn=lambda x: torch.optim.Adam(
x, lr=2e-4, betas=(0.5, 0.999)))
d_B = fe.build(model_fn=Discriminator,
model_name="d_B",
optimizer_fn=lambda x: torch.optim.Adam(
x, lr=2e-4, betas=(0.5, 0.999)))
network = fe.Network(ops=[
ModelOp(inputs="real_A", model=g_AtoB, outputs="fake_B"),
ModelOp(inputs="real_B", model=g_BtoA, outputs="fake_A"),
Buffer(image_in="fake_A", image_out="buffer_fake_A"),
Buffer(image_in="fake_B", image_out="buffer_fake_B"),
ModelOp(inputs="real_A", model=d_A, outputs="d_real_A"),
ModelOp(inputs="fake_A", model=d_A, outputs="d_fake_A"),
ModelOp(inputs="buffer_fake_A", model=d_A, outputs="buffer_d_fake_A"),
ModelOp(inputs="real_B", model=d_B, outputs="d_real_B"),
ModelOp(inputs="fake_B", model=d_B, outputs="d_fake_B"),
ModelOp(inputs="buffer_fake_B", model=d_B, outputs="buffer_d_fake_B"),
ModelOp(inputs="real_A", model=g_BtoA, outputs="same_A"),
ModelOp(inputs="fake_B", model=g_BtoA, outputs="cycled_A"),
ModelOp(inputs="real_B", model=g_AtoB, outputs="same_B"),
ModelOp(inputs="fake_A", model=g_AtoB, outputs="cycled_B"),
GLoss(inputs=("real_A", "d_fake_B", "cycled_A", "same_A"),
weight=weight,
device=device,
outputs="g_AtoB_loss"),
GLoss(inputs=("real_B", "d_fake_A", "cycled_B", "same_B"),
weight=weight,
device=device,
outputs="g_BtoA_loss"),
DLoss(inputs=("d_real_A", "buffer_d_fake_A"),
outputs="d_A_loss",
device=device),
DLoss(inputs=("d_real_B", "buffer_d_fake_B"),
outputs="d_B_loss",
device=device),
UpdateOp(model=g_AtoB, loss_name="g_AtoB_loss"),
UpdateOp(model=g_BtoA, loss_name="g_BtoA_loss"),
UpdateOp(model=d_A, loss_name="d_A_loss"),
UpdateOp(model=d_B, loss_name="d_B_loss")
])
traces = [
ModelSaver(model=g_AtoB, save_dir=save_dir, frequency=10),
ModelSaver(model=g_BtoA, save_dir=save_dir, frequency=10),
LRScheduler(model=g_AtoB, lr_fn=lr_schedule),
LRScheduler(model=g_BtoA, lr_fn=lr_schedule),
LRScheduler(model=d_A, lr_fn=lr_schedule),
LRScheduler(model=d_B, lr_fn=lr_schedule)
]
estimator = fe.Estimator(
network=network,
pipeline=pipeline,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch)
return estimator
def finetune(pretrained_model,
batch_size,
epochs,
model_dir=tempfile.mkdtemp(),
train_steps_per_epoch=None,
eval_steps_per_epoch=None):
train_data, eval_data = cifair10.load_data()
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x",
outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
PadIfNeeded(min_height=40,
min_width=40,
image_in="x",
image_out="x",
mode="train"),
RandomCrop(32,
32,
image_in="x",
image_out="x",
mode="train"),
Sometimes(
HorizontalFlip(image_in="x",
image_out="x",
mode="train")),
CoarseDropout(inputs="x",
outputs="x",
mode="train",
max_holes=1),
ChannelTranspose(inputs="x", outputs="x")
])
model = fe.build(model_fn=lambda: ViTModel(num_classes=100,
image_size=32,
patch_size=4,
num_layers=6,
num_channels=3,
em_dim=256,
num_heads=8,
ff_dim=512),
optimizer_fn=lambda x: torch.optim.SGD(
x, lr=0.01, momentum=0.9, weight_decay=1e-4))
# load the encoder's weight
if hasattr(model, "module"):
model.module.vit_encoder.load_state_dict(
pretrained_model.module.vit_encoder.state_dict())
else:
model.vit_encoder.load_state_dict(
pretrained_model.vit_encoder.state_dict())
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce", from_logits=True),
UpdateOp(model=model, loss_name="ce")
])
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
BestModelSaver(model=model,
save_dir=model_dir,
metric="accuracy",
save_best_mode="max")
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
train_steps_per_epoch=train_steps_per_epoch,
eval_steps_per_epoch=eval_steps_per_epoch)
estimator.fit(warmup=False)
def get_estimator(batch_size=8,
epochs=25,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp(),
data_dir=None):
# load CUB200 dataset.
train_data = cub200.load_data(root_dir=data_dir)
eval_data = train_data.split(0.3)
test_data = eval_data.split(0.5)
# step 1, pipeline
pipeline = fe.Pipeline(batch_size=batch_size,
train_data=train_data,
eval_data=eval_data,
test_data=test_data,
ops=[
ReadImage(inputs="image",
outputs="image",
parent_path=train_data.parent_path),
Normalize(inputs="image",
outputs="image",
mean=1.0,
std=1.0,
max_pixel_value=127.5),
ReadMat(file='annotation',
keys="seg",
parent_path=train_data.parent_path),
Delete(keys="annotation"),
LongestMaxSize(max_size=512,
image_in="image",
image_out="image",
mask_in="seg",
mask_out="seg"),
PadIfNeeded(min_height=512,
min_width=512,
image_in="image",
image_out="image",
mask_in="seg",
mask_out="seg",
border_mode=cv2.BORDER_CONSTANT,
value=0,
mask_value=0),
ShiftScaleRotate(
image_in="image",
mask_in="seg",
image_out="image",
mask_out="seg",
mode="train",
shift_limit=0.2,
rotate_limit=15.0,
scale_limit=0.2,
border_mode=cv2.BORDER_CONSTANT,
value=0,
mask_value=0),
Sometimes(
HorizontalFlip(image_in="image",
mask_in="seg",
image_out="image",
mask_out="seg",
mode="train")),
ChannelTranspose(inputs="image",
outputs="image"),
Reshape(shape=(1, 512, 512),
inputs="seg",
outputs="seg")
])
# step 2, network
resunet50 = fe.build(model_fn=ResUnet50,
model_name="resunet50",
optimizer_fn=lambda x: torch.optim.Adam(x, lr=1e-4))
uncertainty = fe.build(model_fn=UncertaintyLossNet,
model_name="uncertainty",
optimizer_fn=lambda x: torch.optim.Adam(x, lr=1e-5))
network = fe.Network(ops=[
ModelOp(inputs='image',
model=resunet50,
outputs=["label_pred", "mask_pred"]),
CrossEntropy(inputs=["label_pred", "label"],
outputs="cls_loss",
form="sparse",
average_loss=False),
CrossEntropy(inputs=["mask_pred", "seg"],
outputs="seg_loss",
form="binary",
average_loss=False),
ModelOp(inputs=["cls_loss", "seg_loss"],
model=uncertainty,
outputs="total_loss"),
ReduceLoss(inputs="total_loss", outputs="total_loss"),
UpdateOp(model=resunet50, loss_name="total_loss"),
UpdateOp(model=uncertainty, loss_name="total_loss")
])
# step 3, estimator
traces = [
Accuracy(true_key="label", pred_key="label_pred"),
Dice(true_key="seg", pred_key='mask_pred'),
BestModelSaver(model=resunet50,
save_dir=save_dir,
metric="total_loss",
save_best_mode="min"),
LRScheduler(model=resunet50,
lr_fn=lambda step: cosine_decay(
step, cycle_length=13200, init_lr=1e-4))
]
estimator = fe.Estimator(
network=network,
pipeline=pipeline,
traces=traces,
epochs=epochs,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch,
log_steps=500)
return estimator
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 get_estimator(target_size=128,
epochs=55,
save_dir=tempfile.mkdtemp(),
max_train_steps_per_epoch=None,
data_dir=None):
# assert growth parameters
num_grow = np.log2(target_size) - 2
assert num_grow >= 1 and num_grow % 1 == 0, "need exponential of 2 and greater than 8 as target size"
num_phases = int(2 * num_grow + 1)
assert epochs % num_phases == 0, "epoch must be multiple of {} for size {}".format(num_phases, target_size)
num_grow, phase_length = int(num_grow), int(epochs / num_phases)
event_epoch = [1, 1 + phase_length] + [phase_length * (2 * i + 1) + 1 for i in range(1, num_grow)]
event_size = [4] + [2**(i + 3) for i in range(num_grow)]
# set up data schedules
dataset = nih_chestxray.load_data(root_dir=data_dir)
resize_map = {
epoch: Resize(image_in="x", image_out="x", height=size, width=size)
for (epoch, size) in zip(event_epoch, event_size)
}
resize_low_res_map1 = {
epoch: Resize(image_in="x", image_out="x_low_res", height=size // 2, width=size // 2)
for (epoch, size) in zip(event_epoch, event_size)
}
resize_low_res_map2 = {
epoch: Resize(image_in="x_low_res", image_out="x_low_res", height=size, width=size)
for (epoch, size) in zip(event_epoch, event_size)
}
batch_size_map = {
epoch: max(512 // size, 4) * get_num_devices() if size <= 512 else 2 * get_num_devices()
for (epoch, size) in zip(event_epoch, event_size)
}
batch_scheduler = EpochScheduler(epoch_dict=batch_size_map)
pipeline = fe.Pipeline(
batch_size=batch_scheduler,
train_data=dataset,
drop_last=True,
ops=[
ReadImage(inputs="x", outputs="x", color_flag='gray'),
EpochScheduler(epoch_dict=resize_map),
EpochScheduler(epoch_dict=resize_low_res_map1),
EpochScheduler(epoch_dict=resize_low_res_map2),
Normalize(inputs=["x", "x_low_res"], outputs=["x", "x_low_res"], mean=1.0, std=1.0, max_pixel_value=127.5),
ChannelTranspose(inputs=["x", "x_low_res"], outputs=["x", "x_low_res"]),
LambdaOp(fn=lambda: np.random.normal(size=[512]).astype('float32'), outputs="z")
])
fade_in_alpha = torch.tensor(1.0)
d_models = fe.build(
model_fn=lambda: build_D(fade_in_alpha, target_resolution=int(np.log2(target_size)), num_channels=1),
optimizer_fn=[lambda x: Adam(x, lr=0.001, betas=(0.0, 0.99), eps=1e-8)] * len(event_size),
model_name=["d_{}".format(size) for size in event_size])
g_models = fe.build(
model_fn=lambda: build_G(fade_in_alpha, target_resolution=int(np.log2(target_size)), num_channels=1),
optimizer_fn=[lambda x: Adam(x, lr=0.001, betas=(0.0, 0.99), eps=1e-8)] * len(event_size) + [None],
model_name=["g_{}".format(size) for size in event_size] + ["G"])
fake_img_map = {
epoch: ModelOp(inputs="z", outputs="x_fake", model=model)
for (epoch, model) in zip(event_epoch, g_models[:-1])
}
fake_score_map = {
epoch: ModelOp(inputs="x_fake", outputs="fake_score", model=model)
for (epoch, model) in zip(event_epoch, d_models)
}
real_score_map = {
epoch: ModelOp(inputs="x_blend", outputs="real_score", model=model)
for (epoch, model) in zip(event_epoch, d_models)
}
interp_score_map = {
epoch: ModelOp(inputs="x_interp", outputs="interp_score", model=model)
for (epoch, model) in zip(event_epoch, d_models)
}
g_update_map = {
epoch: UpdateOp(loss_name="gloss", model=model)
for (epoch, model) in zip(event_epoch, g_models[:-1])
}
d_update_map = {epoch: UpdateOp(loss_name="dloss", model=model) for (epoch, model) in zip(event_epoch, d_models)}
network = fe.Network(ops=[
EpochScheduler(fake_img_map),
EpochScheduler(fake_score_map),
ImageBlender(alpha=fade_in_alpha, inputs=("x", "x_low_res"), outputs="x_blend"),
EpochScheduler(real_score_map),
Interpolate(inputs=("x_fake", "x"), outputs="x_interp"),
EpochScheduler(interp_score_map),
GradientPenalty(inputs=("x_interp", "interp_score"), outputs="gp"),
GLoss(inputs="fake_score", outputs="gloss"),
DLoss(inputs=("real_score", "fake_score", "gp"), outputs="dloss"),
EpochScheduler(g_update_map),
EpochScheduler(d_update_map)
])
traces = [
AlphaController(alpha=fade_in_alpha,
fade_start_epochs=event_epoch[1:],
duration=phase_length,
batch_scheduler=batch_scheduler,
num_examples=len(dataset)),
ModelSaver(model=g_models[-1], save_dir=save_dir, frequency=phase_length),
ImageSaving(
epoch_model_map={epoch - 1: model
for (epoch, model) in zip(event_epoch[1:] + [epochs + 1], g_models[:-1])},
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)
return estimator
def get_estimator(epochs=20, batch_size=128, max_train_steps_per_epoch=None, save_dir=tempfile.mkdtemp()):
# Dataset Creation
(x_train, y_train), (x_eval, y_eval) = tf.keras.datasets.mnist.load_data()
x_eval0, y_eval0 = x_eval[np.where((y_eval == 1))], np.ones(y_eval[np.where((y_eval == 1))].shape)
x_eval1, y_eval1 = x_eval[np.where((y_eval != 1))], y_eval[np.where((y_eval != 1))]
# Ensuring outliers comprise 50% of the dataset
index = np.random.choice(x_eval1.shape[0], int(x_eval0.shape[0]), replace=False)
x_eval1, y_eval1 = x_eval1[index], np.zeros(y_eval1[index].shape)
x_train, y_train = x_train[np.where((y_train == 1))], np.zeros(y_train[np.where((y_train == 1))].shape)
train_data = fe.dataset.NumpyDataset({"x": x_train, "y": y_train})
x_eval, y_eval = np.concatenate([x_eval0, x_eval1]), np.concatenate([y_eval0, y_eval1])
eval_data = fe.dataset.NumpyDataset({"x": x_eval, "y": y_eval})
pipeline = fe.Pipeline(
train_data=train_data,
eval_data=eval_data,
batch_size=batch_size,
ops=[
ExpandDims(inputs="x", outputs="x"),
Normalize(inputs="x", outputs="x", mean=1.0, std=1.0, max_pixel_value=127.5),
LambdaOp(fn=lambda x: x + np.random.normal(loc=0.0, scale=0.155, size=(28, 28, 1)).astype(np.float32),
inputs="x",
outputs="x_w_noise",
mode="train"),
ChannelTranspose(inputs="x", outputs="x"),
ChannelTranspose(inputs="x_w_noise", outputs="x_w_noise", mode="train")
])
recon_model = fe.build(model_fn=reconstructor,
optimizer_fn=lambda x: torch.optim.RMSprop(x, lr=2e-4),
model_name="reconstructor")
disc_model = fe.build(model_fn=discriminator,
optimizer_fn=lambda x: torch.optim.RMSprop(x, lr=1e-4),
model_name="discriminator")
network = fe.Network(ops=[
ModelOp(model=recon_model, inputs="x_w_noise", outputs="x_fake", mode="train"),
ModelOp(model=recon_model, inputs="x", outputs="x_fake", mode="eval"),
ModelOp(model=disc_model, inputs="x_fake", outputs="fake_score"),
ModelOp(model=disc_model, inputs="x", outputs="true_score"),
RLoss(inputs=("fake_score", "x_fake", "x"), outputs="rloss"),
UpdateOp(model=recon_model, loss_name="rloss"),
DLoss(inputs=("true_score", "fake_score"), outputs="dloss"),
UpdateOp(model=disc_model, loss_name="dloss"),
])
traces = [
F1AUCScores(true_key="y", pred_key="fake_score", mode="eval", output_name=["auc_score", "f1_score"]),
BestModelSaver(model=recon_model, save_dir=save_dir, metric='f1_score', save_best_mode='max'),
BestModelSaver(model=disc_model, save_dir=save_dir, metric='f1_score', save_best_mode='max'),
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
log_steps=50)
return estimator
