def get_callbacks(self, stage: str):
callbacks = {
"scores":
dl.BatchTransformCallback(
input_key="logits",
output_key="scores",
transform=partial(torch.softmax, dim=1),
scope="on_batch_end",
),
"labels":
dl.BatchTransformCallback(
input_key="scores",
output_key="labels",
transform=partial(torch.argmax, dim=1),
scope="on_batch_end",
),
"criterion":
dl.CriterionCallback(metric_key="loss",
input_key="logits",
target_key="targets"),
"optimizer":
dl.OptimizerCallback(
metric_key="loss",
grad_clip_fn=nn.utils.clip_grad_norm_,
grad_clip_params={"max_norm": 1.0},
),
# "scheduler": dl.SchedulerCallback(loader_key="valid", metric_key="loss"),
"accuracy":
dl.AccuracyCallback(input_key="logits",
target_key="targets",
topk_args=(1, 3, 5)),
"classification":
dl.PrecisionRecallF1SupportCallback(input_key="logits",
target_key="targets",
num_classes=10),
"checkpoint":
dl.CheckpointCallback(self._logdir,
loader_key="valid",
metric_key="loss",
minimize=True,
save_n_best=3),
}
if SETTINGS.ml_required:
callbacks["confusion_matrix"] = dl.ConfusionMatrixCallback(
input_key="logits", target_key="targets", num_classes=10)
callbacks["f1_score"] = dl.SklearnBatchCallback(
keys={
"y_pred": "labels",
"y_true": "targets"
},
metric_fn="f1_score",
metric_key="sk_f1",
average="macro",
zero_division=1,
)
return callbacks
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
# sample data
num_samples, num_features, num_classes = int(1e4), int(1e1), 4
X = torch.rand(num_samples, num_features)
y = (torch.rand(num_samples, num_classes) > 0.5).to(torch.float32)
# pytorch loaders
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=32, num_workers=1)
loaders = {"train": loader, "valid": loader}
# model, criterion, optimizer, scheduler
model = torch.nn.Linear(num_features, num_classes)
criterion = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters())
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [2])
# model training
runner = dl.SupervisedRunner(input_key="features",
output_key="logits",
target_key="targets",
loss_key="loss")
callbacks = [
dl.BatchTransformCallback(
transform="F.sigmoid",
scope="on_batch_end",
input_key="logits",
output_key="scores",
),
dl.MultilabelAccuracyCallback(input_key="scores",
target_key="targets",
threshold=0.5),
dl.MultilabelPrecisionRecallF1SupportCallback(
input_key="scores",
target_key="targets",
num_classes=num_classes),
]
if SETTINGS.amp_required and (engine is None or not isinstance(
engine,
(dl.AMPEngine, dl.DataParallelAMPEngine,
dl.DistributedDataParallelAMPEngine),
)):
callbacks.append(
dl.AUCCallback(input_key="scores", target_key="targets"))
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
loaders=loaders,
logdir=logdir,
num_epochs=1,
valid_loader="valid",
valid_metric="accuracy",
minimize_valid_metric=False,
verbose=False,
callbacks=callbacks,
)
if __name__ == "__main__":
args = parser.parse_args()
batch_size = args.batch_size
# 2. model and optimizer
model = get_contrastive_model(args.feature_dim)
optimizer = Adam(model.parameters(), lr=args.learning_rate)
# 3. criterion with triplets sampling
criterion = SupervisedContrastiveLoss(tau=args.temperature)
callbacks = [
dl.BatchTransformCallback(
input_key=["projection_left", "projection_right"],
output_key="full_projection",
scope="on_batch_end",
transform=concat,
),
dl.BatchTransformCallback(
input_key=["target", "target"],
output_key="full_targets",
scope="on_batch_end",
transform=concat,
),
dl.CriterionCallback(
input_key="full_projection", target_key="full_targets", metric_key="loss"
),
]
runner = dl.SelfSupervisedRunner()
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
# sample data
num_users, num_features, num_items = int(1e4), int(1e1), 10
X = torch.rand(num_users, num_features)
y = (torch.rand(num_users, num_items) > 0.5).to(torch.float32)
# pytorch loaders
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=32, num_workers=1)
loaders = {"train": loader, "valid": loader}
# model, criterion, optimizer, scheduler
model = torch.nn.Linear(num_features, num_items)
criterion = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters())
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [2])
callbacks = [
dl.BatchTransformCallback(
input_key="logits",
output_key="scores",
transform=torch.sigmoid,
scope="on_batch_end",
),
dl.CriterionCallback(input_key="logits",
target_key="targets",
metric_key="loss"),
dl.AUCCallback(input_key="scores", target_key="targets"),
dl.HitrateCallback(input_key="scores",
target_key="targets",
topk_args=(1, 3, 5)),
dl.MRRCallback(input_key="scores",
target_key="targets",
topk_args=(1, 3, 5)),
dl.MAPCallback(input_key="scores",
target_key="targets",
topk_args=(1, 3, 5)),
dl.NDCGCallback(input_key="scores",
target_key="targets",
topk_args=(1, 3, 5)),
dl.OptimizerCallback(metric_key="loss"),
dl.SchedulerCallback(),
dl.CheckpointCallback(logdir=logdir,
loader_key="valid",
metric_key="map01",
minimize=False),
]
if SETTINGS.amp_required and (engine is None or not isinstance(
engine,
(dl.AMPEngine, dl.DataParallelAMPEngine,
dl.DistributedDataParallelAMPEngine),
)):
callbacks.append(
dl.AUCCallback(input_key="logits", target_key="targets"))
# model training
runner = dl.SupervisedRunner(input_key="features",
output_key="logits",
target_key="targets",
loss_key="loss")
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
loaders=loaders,
num_epochs=1,
verbose=False,
callbacks=callbacks,
)