eval_batch_size)
ds = tf.data.Dataset.from_tensor_slices((X_train, y_train))
ds_test = tf.data.Dataset.from_tensor_slices((X_test, y_test))
ds_train = prepare(ds,
batch_size,
transform=transform(training=True),
training=True,
buffer_size=n_train)
ds_test = prepare(ds_test,
eval_batch_size,
transform=transform(training=False),
training=False)
setup_runtime()
ds_train, ds_test = distribute_datasets(ds_train, ds_test)
class MixOp(tf.keras.layers.Layer):
def __init__(self):
super().__init__()
self.ops = [
Conv2d(4, 4, kernel_size=1, norm='def', act='def'),
Conv2d(4, 4, kernel_size=3, norm='def', act='def'),
Conv2d(4, 4, kernel_size=3, groups=4, norm='def', act='def'),
Conv2d(4, 4, kernel_size=5, norm='def', act='def'),
Conv2d(4, 4, kernel_size=5, groups=4, norm='def', act='def'),
]
self.alpha = self.add_weight(
if training:
image = cutout(image, 16)
label = tf.one_hot(label, 10)
return image, label
mul = 1
batch_size = 32 * mul
eval_batch_size = batch_size
ds_train, ds_test, steps_per_epoch, test_steps = make_cifar10_dataset(
batch_size, eval_batch_size, transform, sub_ratio=0.01)
setup_runtime(fp16=True)
ds_train, ds_test = distribute_datasets(ds_train, ds_test)
PC_DARTS_cifar = Genotype(normal=[('sep_conv_3x3', 1), ('skip_connect', 0),
('sep_conv_3x3', 0), ('dil_conv_3x3', 1),
('sep_conv_5x5', 0), ('sep_conv_3x3', 1),
('avg_pool_3x3', 0), ('dil_conv_3x3', 1)],
normal_concat=[2, 3, 4, 5],
reduce=[('sep_conv_5x5', 1), ('max_pool_3x3', 0),
('sep_conv_5x5', 1), ('sep_conv_5x5', 2),
('sep_conv_3x3', 0), ('sep_conv_3x3', 3),
('sep_conv_3x3', 1), ('sep_conv_3x3', 2)],
reduce_concat=[2, 3, 4, 5])
drop_path = 0.3
model = NASNet(16, 8, True, drop_path, 10, PC_DARTS_cifar)
def objective(trial: optuna.Trial):
cutout_prob = trial.suggest_float("cutout_prob", 0, 1.0, step=0.1)
mixup_alpha = trial.suggest_float("mixup_alpha", 0, 0.5, step=0.1)
label_smoothing = trial.suggest_uniform("label_smoothing", 0, 0.2)
base_lr = trial.suggest_float("base_lr", 0.01, 0.2, step=0.01)
weight_decay = trial.suggest_loguniform("weight_decay", 1e-5, 1e-3)
ema = trial.suggest_categorical("ema", ["true", "false"])
ema_decay = trial.suggest_loguniform("ema_decay", 0.995,
0.9999) if ema == 'true' else None
@curry
def transform(image, label, training):
if training:
image = random_crop(image, (32, 32), (4, 4))
image = tf.image.random_flip_left_right(image)
image = autoaugment(image, "CIFAR10")
image, label = to_tensor(image, label)
image = normalize(image, [0.491, 0.482, 0.447], [0.247, 0.243, 0.262])
if training:
image = random_apply(cutout(length=16), cutout_prob, image)
label = tf.one_hot(label, 100)
return image, label
def zip_transform(data1, data2):
return mixup(data1, data2, alpha=mixup_alpha)
batch_size = 128
eval_batch_size = 2048
ds_train, ds_test, steps_per_epoch, test_steps = make_cifar100_dataset(
batch_size, eval_batch_size, transform, zip_transform)
setup_runtime(fp16=True)
ds_train, ds_test = distribute_datasets(ds_train, ds_test)
model = ResNet(depth=16, k=8, num_classes=100)
model.build((None, 32, 32, 3))
model.summary()
criterion = CrossEntropy(label_smoothing=label_smoothing)
epochs = 50
lr_schedule = CosineLR(base_lr, steps_per_epoch, epochs=epochs, min_lr=0)
optimizer = SGD(lr_schedule,
momentum=0.9,
weight_decay=weight_decay,
nesterov=True)
train_metrics = {
'loss': Mean(),
'acc': CategoricalAccuracy(),
}
eval_metrics = {
'loss': CategoricalCrossentropy(from_logits=True),
'acc': CategoricalAccuracy(),
}
learner = SuperLearner(model,
criterion,
optimizer,
train_metrics=train_metrics,
eval_metrics=eval_metrics,
work_dir=f"./CIFAR100-NNI",
multiple_steps=True)
callbacks = [OptunaReportIntermediateResult('acc', trial)]
if ema == 'true':
callbacks.append(EMA(ema_decay))
learner.fit(ds_train,
epochs,
ds_test,
val_freq=1,
steps_per_epoch=steps_per_epoch,
val_steps=test_steps,
callbacks=callbacks)
return learner.metric_history.get_metric('acc', "eval")[-1]
