def objective(trial):
logdir = "./logdir"
num_epochs = 10
model = Net(trial)
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
criterion = torch.nn.CrossEntropyLoss()
# model training
runner = SupervisedRunner()
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
loaders=loaders,
logdir=logdir,
num_epochs=num_epochs,
verbose=True,
callbacks=[
AccuracyCallback(),
CatalystPruningCallback(
trial,
metric="accuracy01"), # top-1 accuracy as metric for pruning
],
)
return runner.state.valid_metrics["accuracy01"]
def objective(trial):
lr = trial.suggest_loguniform("lr", 1e-3, 1e-1)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
runner = dl.SupervisedRunner()
runner.train(
model=model,
loaders=loaders,
criterion=criterion,
optimizer=optimizer,
callbacks={
"optuna":
OptunaPruningCallback(loader_key="valid",
metric_key="loss",
minimize=True,
trial=trial),
"accuracy":
AccuracyCallback(num_classes=10,
input_key="logits",
target_key="targets"),
},
num_epochs=2,
valid_metric="accuracy01",
minimize_valid_metric=False,
)
return runner.callbacks["optuna"].best_score
def objective(trial):
lr = trial.suggest_loguniform("lr", 1e-3, 1e-1)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
runner = dl.SupervisedRunner()
runner.train(
model=model,
loaders=loaders,
criterion=criterion,
optimizer=optimizer,
callbacks=[
OptunaCallback(trial),
AccuracyCallback(num_classes=10),
],
num_epochs=10,
main_metric="accuracy01",
minimize_metric=False,
)
return runner.best_valid_metrics[runner.main_metric]
def objective(trial):
logdir = "./logdir"
num_epochs = 10
model = define_model(trial)
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
criterion = torch.nn.CrossEntropyLoss()
# model training
runner = SupervisedRunner()
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
loaders=loaders,
logdir=logdir,
num_epochs=num_epochs,
verbose=True,
callbacks={
# NOTE(crcrpar): Consult above [Why OptunaPruningCallback] for the use of
# Catalys's callback for Optuna, not Optuna's one for Catalyst.
# top-1 accuracy as metric for pruning
"optuna":
OptunaPruningCallback(
loader_key="valid",
metric_key="accuracy01",
minimize=False,
trial=trial,
),
"accuracy":
AccuracyCallback(
input_key="logits",
target_key="targets",
num_classes=10,
),
},
)
return runner.callbacks["optuna"].best_score
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--seed', type=int, default=42, help='Random seed')
parser.add_argument('-dd',
'--data-dir',
type=str,
default='data',
help='Data directory')
parser.add_argument('-l',
'--loss',
type=str,
default='label_smooth_cross_entropy')
parser.add_argument('-t1', '--temper1', type=float, default=0.2)
parser.add_argument('-t2', '--temper2', type=float, default=4.0)
parser.add_argument('-optim', '--optimizer', type=str, default='adam')
parser.add_argument('-prep', '--prep_function', type=str, default='none')
parser.add_argument('--train_on_different_datasets', action='store_true')
parser.add_argument('--use-current', action='store_true')
parser.add_argument('--use-extra', action='store_true')
parser.add_argument('--use-unlabeled', action='store_true')
parser.add_argument('--fast', action='store_true')
parser.add_argument('--mixup', action='store_true')
parser.add_argument('--balance', action='store_true')
parser.add_argument('--balance-datasets', action='store_true')
parser.add_argument('--show', action='store_true')
parser.add_argument('-v', '--verbose', action='store_true')
parser.add_argument('-m',
'--model',
type=str,
default='efficientnet-b4',
help='')
parser.add_argument('-b',
'--batch-size',
type=int,
default=8,
help='Batch Size during training, e.g. -b 64')
parser.add_argument('-e',
'--epochs',
type=int,
default=100,
help='Epoch to run')
parser.add_argument('-s',
'--sizes',
default=380,
type=int,
help='Image size for training & inference')
parser.add_argument('-f', '--fold', type=int, default=None)
parser.add_argument('-t', '--transfer', default=None, type=str, help='')
parser.add_argument('-lr',
'--learning_rate',
type=float,
default=1e-4,
help='Initial learning rate')
parser.add_argument('-a',
'--augmentations',
default='medium',
type=str,
help='')
parser.add_argument('-accum', '--accum-step', type=int, default=1)
parser.add_argument('-metric', '--metric', type=str, default='accuracy01')
args = parser.parse_args()
diff_dataset_train = args.train_on_different_datasets
data_dir = args.data_dir
epochs = args.epochs
batch_size = args.batch_size
seed = args.seed
loss_name = args.loss
optim_name = args.optimizer
prep_function = args.prep_function
model_name = args.model
size = args.sizes,
print(size)
print(size[0])
image_size = (size[0], size[0])
print(image_size)
fast = args.fast
fold = args.fold
mixup = args.mixup
balance = args.balance
balance_datasets = args.balance_datasets
show_batches = args.show
verbose = args.verbose
use_current = args.use_current
use_extra = args.use_extra
use_unlabeled = args.use_unlabeled
learning_rate = args.learning_rate
augmentations = args.augmentations
transfer = args.transfer
accum_step = args.accum_step
#cosine_loss accuracy01
main_metric = args.metric
print(data_dir)
num_classes = 5
assert use_current or use_extra
print(fold)
current_time = datetime.now().strftime('%b%d_%H_%M')
random_name = get_random_name()
current_time = datetime.now().strftime('%b%d_%H_%M')
random_name = get_random_name()
# if folds is None or len(folds) == 0:
# folds = [None]
torch.cuda.empty_cache()
checkpoint_prefix = f'{model_name}_{size}_{augmentations}'
if transfer is not None:
checkpoint_prefix += '_pretrain_from_' + str(transfer)
else:
if use_current:
checkpoint_prefix += '_current'
if use_extra:
checkpoint_prefix += '_extra'
if use_unlabeled:
checkpoint_prefix += '_unlabeled'
if fold is not None:
checkpoint_prefix += f'_fold{fold}'
directory_prefix = f'{current_time}_{checkpoint_prefix}'
log_dir = os.path.join('runs', directory_prefix)
os.makedirs(log_dir, exist_ok=False)
set_manual_seed(seed)
model = get_model(model_name)
if transfer is not None:
print("Transfering weights from model checkpoint")
model.load_state_dict(torch.load(transfer)['model_state_dict'])
model = model.cuda()
if diff_dataset_train:
train_on = ['current_train', 'extra_train']
valid_on = ['unlabeled']
train_ds, valid_ds, train_sizes = get_datasets_universal(
train_on=train_on,
valid_on=valid_on,
image_size=image_size,
augmentation=augmentations,
target_dtype=int,
prep_function=prep_function)
else:
train_ds, valid_ds, train_sizes = get_datasets(
data_dir=data_dir,
use_current=use_current,
use_extra=use_extra,
image_size=image_size,
prep_function=prep_function,
augmentation=augmentations,
target_dtype=int,
fold=fold,
folds=5)
train_loader, valid_loader = get_dataloaders(train_ds,
valid_ds,
batch_size=batch_size,
train_sizes=train_sizes,
num_workers=6,
balance=True,
balance_datasets=True,
balance_unlabeled=False)
loaders = collections.OrderedDict()
loaders["train"] = train_loader
loaders["valid"] = valid_loader
runner = SupervisedRunner(input_key='image')
criterions = get_loss(loss_name)
# criterions_tempered = TemperedLogLoss()
# optimizer = catalyst.contrib.nn.optimizers.radam.RAdam(model.parameters(), lr = learning_rate)
optimizer = get_optim(optim_name, model, learning_rate)
# optimizer = catalyst.contrib.nn.optimizers.Adam(model.parameters(), lr = learning_rate)
# criterions = nn.CrossEntropyLoss()
# optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[25], gamma=0.8)
# cappa = CappaScoreCallback()
Q = math.floor(len(train_ds) / batch_size)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=Q)
if main_metric != 'accuracy01':
callbacks = [
AccuracyCallback(num_classes=num_classes),
CosineLossCallback(),
OptimizerCallback(accumulation_steps=accum_step),
CheckpointCallback(save_n_best=epochs)
]
else:
callbacks = [
AccuracyCallback(num_classes=num_classes),
OptimizerCallback(accumulation_steps=accum_step),
CheckpointCallback(save_n_best=epochs)
]
# main_metric = 'accuracy01'
runner.train(
fp16=True,
model=model,
criterion=criterions,
optimizer=optimizer,
scheduler=scheduler,
callbacks=callbacks,
loaders=loaders,
logdir=log_dir,
num_epochs=epochs,
verbose=verbose,
main_metric=main_metric,
minimize_metric=False,
)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[3, 8],
gamma=0.3)
# model runner
runner = SupervisedRunner()
# model training
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
loaders=loaders,
callbacks=[
AccuracyCallback(topk_args=[1, 3, 5]),
EarlyStoppingCallback(patience=2, min_delta=0.01),
],
logdir=logdir,
num_epochs=num_epochs,
check=True,
)
# In[ ]:
# utils.plot_metrics(
# logdir=logdir,
# metrics=["loss", "accuracy01", "accuracy03", "_base/lr"])
# # Setup 5 - training with 1cycle
def train(self):
# model = {"model": self.model}
# criterion = {"criterion": nn.CrossEntropyLoss()}
# optimizer = {"optimizer": self.optimizer}
callbacks = [
# dl.CriterionCallback(
# input_key="logits",
# target_key="targets",
# metric_key="loss",
# criterion_key="criterion",
# ),
# dl.OptimizerCallback(
# model_key="model",
# optimizer_key="optimizer",
# metric_key="loss"
# ),
EarlyStoppingCallback(patience=15,
metric_key="loss",
loader_key="valid",
minimize=True,
min_delta=0),
AccuracyCallback(num_classes=2,
input_key="logits",
target_key="targets"),
AUCCallback(input_key="logits", target_key="targets"),
# CheckpointCallback(
# "./logs", loader_key="valid", metric_key="loss", minimize=True, save_n_best=3,
# # load_on_stage_start={"model": "best"},
# load_on_stage_end={"model": "best"}
# ),
]
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,
mode="min")
train_dataset = TensorDataset(self.tr_eps, self.tr_labels)
val_dataset = TensorDataset(self.val_eps, self.val_labels)
test_dataset = TensorDataset(self.tst_eps, self.test_labels)
runner = CustomRunner("./logs")
v_bs = self.val_eps.shape[0]
t_bs = self.tst_eps.shape[0]
loaders = {
"train":
DataLoader(
train_dataset,
batch_size=self.batch_size,
num_workers=0,
shuffle=True,
),
"valid":
DataLoader(
val_dataset,
batch_size=v_bs,
num_workers=0,
shuffle=True,
),
}
if self.complete_arc == True:
if self.PT in ["milc", "two-loss-milc"]:
if self.exp in ["UFPT", "FPT"]:
model_dict = torch.load(
os.path.join(self.oldpath, "best_full" + ".pth"),
map_location=self.device,
)
model_dict = model_dict["model_state_dict"]
print("Complete Arch Loaded")
self.model.load_state_dict(model_dict)
# num_features=2
# model training
# train_loader_param = {"batch_size": 64,
# "shuffle":True,
# }
# val_loader_param = {"batch_size": 32,
# "shuffle": True,
# }
# loaders_params = {"train" : train_loader_param,
# "valid": val_loader_param}
# datasets = {
# "batch_size": 64,
# "num_workers": 1,
# "loaders_params": loaders_params,
# "get_datasets_fn": self.datasets_fn,
# "num_features": num_features,
# },
runner.train(
model=self.model,
optimizer=self.optimizer,
# criterion=criterion,
scheduler=scheduler,
loaders=loaders,
valid_loader='valid',
callbacks=callbacks,
logdir="./logs",
num_epochs=self.epochs,
verbose=True,
load_best_on_end=True,
valid_metric="loss",
minimize_valid_metric=True,
)
loader = (DataLoader(test_dataset,
batch_size=t_bs,
num_workers=1,
shuffle=True), )
(
self.test_accuracy,
self.test_auc,
self.test_loss,
) = runner.predict_batch(next(iter(loader)))
def get_criterions(
modification_flag,
modification_type,
embedding_loss,
mask_loss,
bits_loss,
num_epochs: int,
feature_maps_loss=None,
mixup=False,
cutmix=False,
tsa=False,
class_names=["Cover", "JMiPOD", "JUNIWARD", "UERD"],
):
criterions_dict = {}
callbacks = []
losses = []
need_embedding_auc_score = False
if modification_flag is not None:
for criterion in modification_flag:
if isinstance(criterion, (list, tuple)):
loss_name, loss_weight = criterion
else:
loss_name, loss_weight = criterion, 1.0
cd, criterion, criterion_name = get_criterion_callback(
loss_name,
num_epochs=num_epochs,
input_key=INPUT_TRUE_MODIFICATION_FLAG,
output_key=OUTPUT_PRED_MODIFICATION_FLAG,
prefix=f"modification_flag/{loss_name}",
loss_weight=float(loss_weight),
mixup=mixup,
cutmix=cutmix,
tsa=tsa,
)
criterions_dict.update(cd)
callbacks.append(criterion)
losses.append(criterion_name)
print("Using loss", loss_name, loss_weight)
# Metrics
callbacks += [
ConfusionMatrixCallback(
input_key=INPUT_TRUE_MODIFICATION_FLAG,
output_key=OUTPUT_PRED_MODIFICATION_FLAG,
prefix="flag",
activation_fn=lambda x: x.sigmoid() > 0.5,
class_names=["Original", "Modified"],
),
CompetitionMetricCallback(
input_key=INPUT_TRUE_MODIFICATION_TYPE,
output_key=OUTPUT_PRED_MODIFICATION_FLAG,
prefix="auc",
output_activation=binary_logits_to_probas,
class_names=class_names,
),
OutputDistributionCallback(
input_key=INPUT_TRUE_MODIFICATION_FLAG,
output_key=OUTPUT_PRED_MODIFICATION_FLAG,
output_activation=binary_logits_to_probas,
prefix="distribution/binary",
),
AccuracyCallback(
prefix="accuracy_b",
activation="Sigmoid",
input_key=INPUT_TRUE_MODIFICATION_FLAG,
output_key=OUTPUT_PRED_MODIFICATION_FLAG,
accuracy_args=[1],
num_classes=1,
),
BestMetricCheckpointCallback(target_metric="auc", target_metric_minimize=False, save_n_best=3),
]
if modification_type is not None:
for criterion in modification_type:
if isinstance(criterion, (list, tuple)):
loss_name, loss_weight = criterion
else:
loss_name, loss_weight = criterion, 1.0
cd, criterion, criterion_name = get_criterion_callback(
loss_name,
num_epochs=num_epochs,
input_key=INPUT_TRUE_MODIFICATION_TYPE,
output_key=OUTPUT_PRED_MODIFICATION_TYPE,
prefix=f"modification_type/{loss_name}",
loss_weight=float(loss_weight),
mixup=mixup,
cutmix=cutmix,
tsa=tsa,
)
criterions_dict.update(cd)
callbacks.append(criterion)
losses.append(criterion_name)
print("Using loss", loss_name, loss_weight)
# Metrics
callbacks += [
ConfusionMatrixCallback(
input_key=INPUT_TRUE_MODIFICATION_TYPE,
output_key=OUTPUT_PRED_MODIFICATION_TYPE,
prefix="type",
class_names=class_names,
),
AccuracyCallback(
input_key=INPUT_TRUE_MODIFICATION_TYPE, output_key=OUTPUT_PRED_MODIFICATION_TYPE, prefix="accuracy"
),
CompetitionMetricCallback(
input_key=INPUT_TRUE_MODIFICATION_TYPE,
output_key=OUTPUT_PRED_MODIFICATION_TYPE,
output_activation=classifier_logits_to_probas,
prefix="auc_classifier",
class_names=class_names,
),
OutputDistributionCallback(
input_key=INPUT_TRUE_MODIFICATION_FLAG,
output_key=OUTPUT_PRED_MODIFICATION_TYPE,
output_activation=classifier_logits_to_probas,
prefix="distribution/classifier",
),
BestMetricCheckpointCallback(target_metric="auc_classifier", target_metric_minimize=False, save_n_best=3),
]
if mask_loss is not None:
for criterion in mask_loss:
if isinstance(criterion, (list, tuple)):
loss_name, loss_weight = criterion
else:
loss_name, loss_weight = criterion, 1.0
cd, criterion, criterion_name = get_criterion_callback(
loss_name,
num_epochs=num_epochs,
input_key=INPUT_TRUE_MODIFICATION_MASK,
output_key=OUTPUT_PRED_MODIFICATION_MASK,
prefix=f"mask/{loss_name}",
loss_weight=float(loss_weight),
mixup=mixup,
cutmix=cutmix,
tsa=tsa,
)
criterions_dict.update(cd)
callbacks.append(criterion)
losses.append(criterion_name)
print("Using loss", loss_name, loss_weight)
# Metrics
callbacks += [
CompetitionMetricCallbackFromMask(
input_key=INPUT_TRUE_MODIFICATION_MASK, output_key=OUTPUT_PRED_MODIFICATION_MASK, prefix="auc_mask"
),
BestMetricCheckpointCallback(target_metric="auc_mask", target_metric_minimize=False, save_n_best=3),
]
if embedding_loss is not None:
for criterion in embedding_loss:
if isinstance(criterion, (list, tuple)):
loss_name, loss_weight = criterion
else:
loss_name, loss_weight = criterion, 1.0
cd, criterion, criterion_name = get_criterion_callback(
loss_name,
num_epochs=num_epochs,
input_key=INPUT_TRUE_MODIFICATION_TYPE,
output_key=OUTPUT_PRED_EMBEDDING_ARC_MARGIN,
prefix=f"embedding/{loss_name}",
loss_weight=float(loss_weight),
mixup=mixup,
cutmix=cutmix,
tsa=tsa,
)
criterions_dict.update(cd)
callbacks.append(criterion)
losses.append(criterion_name)
print("Using loss", loss_name, loss_weight)
if loss_name == "arc_face":
need_embedding_auc_score = True
if need_embedding_auc_score:
callbacks += [
OutputDistributionCallback(
input_key=INPUT_TRUE_MODIFICATION_FLAG,
output_key=OUTPUT_PRED_EMBEDDING_ARC_MARGIN,
output_activation=classifier_logits_to_probas,
prefix="distribution/embedding",
),
CompetitionMetricCallback(
input_key=INPUT_TRUE_MODIFICATION_TYPE,
output_key=OUTPUT_PRED_EMBEDDING_ARC_MARGIN,
prefix="auc_embedding",
output_activation=classifier_logits_to_probas,
class_names=class_names,
),
BestMetricCheckpointCallback(
target_metric="auc_embedding", target_metric_minimize=False, save_n_best=3
),
]
if feature_maps_loss is not None:
for criterion in feature_maps_loss:
if isinstance(criterion, (list, tuple)):
loss_name, loss_weight = criterion
else:
loss_name, loss_weight = criterion, 1.0
for fm in {OUTPUT_FEATURE_MAP_4, OUTPUT_FEATURE_MAP_8, OUTPUT_FEATURE_MAP_16, OUTPUT_FEATURE_MAP_32}:
cd, criterion, criterion_name = get_criterion_callback(
loss_name,
num_epochs=num_epochs,
input_key=INPUT_TRUE_MODIFICATION_TYPE,
output_key=fm,
prefix=f"{fm}/{loss_name}",
loss_weight=float(loss_weight),
mixup=mixup,
cutmix=cutmix,
tsa=tsa,
)
criterions_dict.update(cd)
callbacks.append(criterion)
losses.append(criterion_name)
print("Using loss", fm, loss_name, loss_weight)
if bits_loss is not None:
for criterion in bits_loss:
if isinstance(criterion, (list, tuple)):
loss_name, loss_weight = criterion
else:
loss_name, loss_weight = criterion, 1.0
cd, criterion, criterion_name = get_criterion_callback(
loss_name,
num_epochs=num_epochs,
input_key=INPUT_TRUE_PAYLOAD_BITS,
output_key=OUTPUT_PRED_PAYLOAD_BITS,
prefix=f"bits/{loss_name}",
loss_weight=float(loss_weight),
mixup=mixup,
cutmix=cutmix,
tsa=tsa,
)
criterions_dict.update(cd)
callbacks.append(criterion)
losses.append(criterion_name)
print("Using loss", loss_name, loss_weight)
# Metrics
callbacks += [
CompetitionMetricCallback(
input_key=INPUT_TRUE_MODIFICATION_TYPE,
output_key=OUTPUT_PRED_PAYLOAD_BITS,
output_activation=log_plus_one,
prefix="auc_bits",
),
BestMetricCheckpointCallback(target_metric="auc_bits", target_metric_minimize=False, save_n_best=3),
]
callbacks.append(CriterionAggregatorCallback(prefix="loss", loss_keys=losses))
if mixup:
callbacks.append(MixupInputCallback(fields=[INPUT_IMAGE_KEY], alpha=0.5, p=0.5))
return criterions_dict, callbacks
