def test_disabling_loss_for_train():
old_stdout = sys.stdout
sys.stdout = str_stdout = StringIO()
# experiment_setup
logdir = "./logs/control_flow"
checkpoint = logdir + "/checkpoints"
logfile = checkpoint + "/_metrics.json"
# data
num_samples, num_features = int(1e4), int(1e1)
X = torch.rand(num_samples, num_features)
y = torch.randint(0, 5, size=[num_samples])
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, 5)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
runner = dl.SupervisedRunner()
n_epochs = 5
# first stage
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
loaders=loaders,
logdir=logdir,
num_epochs=n_epochs,
verbose=False,
main_metric="accuracy01",
callbacks=[
dl.ControlFlowCallback(dl.CriterionCallback(),
ignore_loaders=["train"]),
dl.AccuracyCallback(accuracy_args=[1, 3, 5]),
dl.CheckRunCallback(num_epoch_steps=n_epochs),
],
)
sys.stdout = old_stdout
exp_output = str_stdout.getvalue()
assert len(re.findall(r"\(train\): loss", exp_output)) == 5
assert len(re.findall(r"\(valid\): loss", exp_output)) == 0
assert len(re.findall(r".*/train\.\d\.pth", exp_output)) == 1
assert os.path.isfile(logfile)
assert os.path.isfile(checkpoint + "/best.pth")
assert os.path.isfile(checkpoint + "/best_full.pth")
assert os.path.isfile(checkpoint + "/last.pth")
assert os.path.isfile(checkpoint + "/last_full.pth")
pth_files = [
file for file in os.listdir(checkpoint) if file.endswith(".pth")
]
assert len(pth_files) == 6
shutil.rmtree(logdir, ignore_errors=True)
def test_metric_learning_pipeline():
"""
Test if classification pipeline can run and compute metrics.
In this test we check that LoaderMetricCallback works with
CMCMetric (ICallbackLoaderMetric).
"""
with TemporaryDirectory() as tmp_dir:
dataset_train = datasets.MnistMLDataset(root=tmp_dir, download=True)
sampler = data.BalanceBatchSampler(labels=dataset_train.get_labels(), p=5, k=10)
train_loader = DataLoader(
dataset=dataset_train, sampler=sampler, batch_size=sampler.batch_size,
)
dataset_val = datasets.MnistQGDataset(root=tmp_dir, transform=None, gallery_fraq=0.2)
val_loader = DataLoader(dataset=dataset_val, batch_size=1024)
model = DummyModel(num_features=28 * 28, num_classes=NUM_CLASSES)
optimizer = Adam(model.parameters(), lr=0.001)
sampler_inbatch = data.HardTripletsSampler(norm_required=False)
criterion = nn.TripletMarginLossWithSampler(margin=0.5, sampler_inbatch=sampler_inbatch)
callbacks = OrderedDict(
{
"cmc": dl.ControlFlowCallback(
LoaderMetricCallback(
CMCMetric(
topk_args=[1],
embeddings_key="embeddings",
labels_key="targets",
is_query_key="is_query",
),
input_key=["embeddings", "is_query"],
target_key=["targets"],
),
loaders="valid",
),
"control": dl.PeriodicLoaderCallback(
valid_loader_key="valid", valid_metric_key="cmc", valid=2
),
}
)
runner = CustomRunner(input_key="features", output_key="embeddings")
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
callbacks=callbacks,
loaders=OrderedDict({"train": train_loader, "valid": val_loader}),
verbose=False,
valid_loader="valid",
num_epochs=4,
)
assert "cmc01" in runner.loader_metrics
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
from catalyst import utils
utils.set_global_seed(RANDOM_STATE)
# 1. train, valid and test loaders
transforms = Compose([ToTensor(), Normalize((0.1307, ), (0.3081, ))])
train_data = MNIST(os.getcwd(),
train=True,
download=True,
transform=transforms)
train_labels = train_data.targets.cpu().numpy().tolist()
train_sampler = data.BatchBalanceClassSampler(train_labels,
num_classes=10,
num_samples=4)
train_loader = DataLoader(train_data, batch_sampler=train_sampler)
valid_dataset = MNIST(root=os.getcwd(),
transform=transforms,
train=False,
download=True)
valid_loader = DataLoader(dataset=valid_dataset, batch_size=32)
test_dataset = MNIST(root=os.getcwd(),
transform=transforms,
train=False,
download=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=32)
# 2. model and optimizer
model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 16),
nn.LeakyReLU(inplace=True))
optimizer = Adam(model.parameters(), lr=LR)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [2])
# 3. criterion with triplets sampling
sampler_inbatch = data.HardTripletsSampler(norm_required=False)
criterion = nn.TripletMarginLossWithSampler(
margin=0.5, sampler_inbatch=sampler_inbatch)
# 4. training with catalyst Runner
class CustomRunner(dl.SupervisedRunner):
def handle_batch(self, batch) -> None:
images, targets = batch["features"].float(
), batch["targets"].long()
features = self.model(images)
self.batch = {
"embeddings": features,
"targets": targets,
}
callbacks = [
dl.ControlFlowCallback(
dl.CriterionCallback(input_key="embeddings",
target_key="targets",
metric_key="loss"),
loaders="train",
),
dl.SklearnModelCallback(
feature_key="embeddings",
target_key="targets",
train_loader="train",
valid_loaders=["valid", "infer"],
model_fn=RandomForestClassifier,
predict_method="predict_proba",
predict_key="sklearn_predict",
random_state=RANDOM_STATE,
n_estimators=50,
),
dl.ControlFlowCallback(
dl.AccuracyCallback(target_key="targets",
input_key="sklearn_predict",
topk_args=(1, 3)),
loaders=["valid", "infer"],
),
]
runner = CustomRunner(input_key="features", output_key="embeddings")
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
callbacks=callbacks,
loaders={
"train": train_loader,
"valid": valid_loader,
"infer": test_loader
},
verbose=False,
valid_loader="valid",
valid_metric="accuracy",
minimize_valid_metric=False,
num_epochs=TRAIN_EPOCH,
logdir=logdir,
)
valid_path = Path(logdir) / "logs/infer.csv"
best_accuracy = max(
float(row["accuracy"]) for row in read_csv(valid_path))
assert best_accuracy > 0.8
pair_transform=False),
download=True,
)
train_loader = DataLoader(train_data,
batch_size=batch_size,
shuffle=True,
pin_memory=True)
valid_loader = DataLoader(test_data,
batch_size=batch_size,
pin_memory=True)
callbacks = [
dl.ControlFlowCallback(
dl.CriterionCallback(input_key="out_1",
target_key="out_2",
metric_key="loss"),
loaders="train",
),
dl.SklearnModelCallback(
feature_key="embeddings",
target_key="targets",
train_loader="train",
valid_loaders="valid",
model_fn=LogisticRegression,
predict_key="sklearn_predict",
predict_method="predict_proba",
),
dl.OptimizerCallback(metric_key="loss"),
dl.ControlFlowCallback(
dl.AccuracyCallback(target_key="targets",
input_key="sklearn_predict",
def main(args):
wandb.init(project="teacher-pruning", config=vars(args))
set_global_seed(42)
# dataloader initialization
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
train_dataset = Wrp(
datasets.CIFAR10(root=os.getcwd(),
train=True,
transform=transform_train,
download=True))
valid_dataset = Wrp(
datasets.CIFAR10(root=os.getcwd(),
train=False,
transform=transform_test))
train_dataloader = DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=2)
valid_dataloader = DataLoader(dataset=valid_dataset,
batch_size=128,
num_workers=2)
loaders = {
"train": train_dataloader,
"valid": valid_dataloader,
}
# model initialization
model = PreActResNet18()
model.fc = nn.Linear(512, 10)
if args.teacher_model is not None:
is_kd = True
teacher_model = NAME2MODEL[args.teacher_model]()
load_model_from_path(model=teacher_model, path=args.teacher_path)
model = {
"student": model,
"teacher": teacher_model,
}
output_hiddens = args.beta is None
is_kd_on_hiddens = output_hiddens
runner = KDRunner(device=args.device, output_hiddens=output_hiddens)
parameters = model["student"].parameters()
else:
is_kd = False
runner = dl.SupervisedRunner(device=args.device)
parameters = model.parameters()
# optimizer
optimizer_cls = NAME2OPTIM[args.optimizer]
optimizer_kwargs = {"params": parameters, "lr": args.lr}
if args.optimizer == "sgd":
optimizer_kwargs["momentum"] = args.momentum
else:
optimizer_kwargs["betas"] = (args.beta1, args.beta2)
optimizer = optimizer_cls(**optimizer_kwargs)
scheduler = MultiStepLR(optimizer, milestones=[80, 120], gamma=args.gamma)
logdir = f"logs/{wandb.run.name}"
# callbacks
callbacks = [dl.AccuracyCallback(num_classes=10), WandbCallback()]
if is_kd:
metrics = {}
callbacks.append(dl.CriterionCallback(output_key="cls_loss"))
callbacks.append(DiffOutputCallback())
coefs = get_loss_coefs(args.alpha, args.beta)
metrics["cls_loss"] = coefs[0]
metrics["diff_output_loss"] = coefs[1]
if is_kd_on_hiddens:
callbacks.append(DiffHiddenCallback())
metrics["diff_hidden_loss"] = coefs[2]
aggregator_callback = dl.MetricAggregationCallback(prefix="loss",
metrics=metrics,
mode="weighted_sum")
wrapped_agg_callback = dl.ControlFlowCallback(aggregator_callback,
loaders=["train"])
callbacks.append(wrapped_agg_callback)
runner.train(
model=model,
optimizer=optimizer,
scheduler=scheduler,
criterion=nn.CrossEntropyLoss(),
loaders=loaders,
callbacks=callbacks,
num_epochs=args.epoch,
logdir=logdir,
verbose=True,
)
def run_ml_pipeline(sampler_inbatch: data.IInbatchTripletSampler) -> float:
"""
Full metric learning pipeline, including train and val.
This function is also used as minimal example in README.md, section name:
'CV - MNIST with Metric Learning'.
Args:
sampler_inbatch: sampler to forming triplets
Returns:
best metric value
"""
# 1. train and valid datasets
dataset_root = "./data"
transforms = t.Compose([t.ToTensor(), t.Normalize((0.1307, ), (0.3081, ))])
dataset_train = datasets.MnistMLDataset(
root=dataset_root,
train=True,
download=True,
transform=transforms,
)
sampler = data.BalanceBatchSampler(labels=dataset_train.get_labels(),
p=5,
k=10)
train_loader = DataLoader(dataset=dataset_train,
sampler=sampler,
batch_size=sampler.batch_size)
dataset_val = datasets.MnistQGDataset(root=dataset_root,
transform=transforms,
gallery_fraq=0.2)
val_loader = DataLoader(dataset=dataset_val, batch_size=1024)
# 2. model and optimizer
model = models.SimpleConv(features_dim=16)
optimizer = Adam(model.parameters(), lr=0.0005)
# 3. criterion with triplets sampling
criterion = nn.TripletMarginLossWithSampler(
margin=0.5, sampler_inbatch=sampler_inbatch)
# 4. training with catalyst Runner
callbacks = [
dl.ControlFlowCallback(dl.CriterionCallback(), loaders="train"),
dl.ControlFlowCallback(dl.CMCScoreCallback(topk_args=[1]),
loaders="valid"),
dl.PeriodicLoaderCallback(valid=100),
]
runner = dl.SupervisedRunner(device=utils.get_device())
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
callbacks=callbacks,
loaders={
"train": train_loader,
"valid": val_loader
},
minimize_metric=False,
verbose=True,
valid_loader="valid",
num_epochs=100,
main_metric="cmc01",
)
return runner.best_valid_metrics["cmc01"]
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
from catalyst import utils
utils.set_global_seed(RANDOM_STATE)
# 1. generate data
num_samples, num_features, num_classes = int(1e4), int(30), 3
X, y = make_classification(
n_samples=num_samples,
n_features=num_features,
n_informative=num_features,
n_repeated=0,
n_redundant=0,
n_classes=num_classes,
n_clusters_per_class=1,
)
X, y = torch.tensor(X), torch.tensor(y)
dataset = TensorDataset(X, y)
loader = DataLoader(dataset,
batch_size=64,
num_workers=1,
shuffle=True)
# 2. model, optimizer and scheduler
hidden_size, out_features = 20, 16
model = nn.Sequential(nn.Linear(num_features, hidden_size), nn.ReLU(),
nn.Linear(hidden_size, out_features))
optimizer = Adam(model.parameters(), lr=LR)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [2])
# 3. criterion with triplets sampling
sampler_inbatch = data.HardTripletsSampler(norm_required=False)
criterion = nn.TripletMarginLossWithSampler(
margin=0.5, sampler_inbatch=sampler_inbatch)
# 4. training with catalyst Runner
class CustomRunner(dl.SupervisedRunner):
def handle_batch(self, batch) -> None:
features, targets = batch["features"].float(
), batch["targets"].long()
embeddings = self.model(features)
self.batch = {
"embeddings": embeddings,
"targets": targets,
}
callbacks = [
dl.SklearnModelCallback(
feature_key="embeddings",
target_key="targets",
train_loader="train",
valid_loaders="valid",
model_fn=RandomForestClassifier,
predict_method="predict_proba",
predict_key="sklearn_predict",
random_state=RANDOM_STATE,
n_estimators=100,
),
dl.ControlFlowCallback(
dl.AccuracyCallback(target_key="targets",
input_key="sklearn_predict",
topk_args=(1, 3)),
loaders="valid",
),
]
runner = CustomRunner(input_key="features", output_key="embeddings")
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
callbacks=callbacks,
scheduler=scheduler,
loaders={
"train": loader,
"valid": loader
},
verbose=False,
valid_loader="valid",
valid_metric="accuracy",
minimize_valid_metric=False,
num_epochs=TRAIN_EPOCH,
logdir=logdir,
)
valid_path = Path(logdir) / "logs/valid.csv"
best_accuracy = max(
float(row["accuracy"]) for row in read_csv(valid_path))
assert best_accuracy > 0.9
dl.CriterionCallback(input_key="projection_left",
target_key="projection_right",
metric_key="loss"),
dl.SklearnModelCallback(
feature_key="embedding_origin",
target_key="target",
train_loader="train",
valid_loaders="valid",
model_fn=LogisticRegression,
predict_key="sklearn_predict",
predict_method="predict_proba",
),
dl.OptimizerCallback(metric_key="loss"),
dl.ControlFlowCallback(
dl.AccuracyCallback(target_key="target",
input_key="sklearn_predict",
topk_args=(1, 3)),
loaders="valid",
),
]
runner = SelfSupervisedRunner()
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
callbacks=callbacks,
loaders=get_loaders(args.dataset, args.batch_size, args.num_workers),
verbose=True,
logdir=args.logdir,
valid_loader="train",
def test_reid_pipeline():
"""This test checks that reid pipeline runs and compute metrics with ReidCMCScoreCallback"""
with TemporaryDirectory() as logdir:
# 1. train and valid loaders
transforms = Compose([ToTensor(), Normalize((0.1307, ), (0.3081, ))])
train_dataset = MnistMLDataset(root=os.getcwd(),
download=True,
transform=transforms)
sampler = data.BatchBalanceClassSampler(
labels=train_dataset.get_labels(),
num_classes=3,
num_samples=10,
num_batches=20)
train_loader = DataLoader(dataset=train_dataset,
batch_sampler=sampler,
num_workers=0)
valid_dataset = MnistReIDQGDataset(root=os.getcwd(),
transform=transforms,
gallery_fraq=0.2)
valid_loader = DataLoader(dataset=valid_dataset, batch_size=1024)
# 2. model and optimizer
model = models.MnistSimpleNet(out_features=16)
optimizer = Adam(model.parameters(), lr=0.001)
# 3. criterion with triplets sampling
sampler_inbatch = data.AllTripletsSampler(max_output_triplets=1000)
criterion = nn.TripletMarginLossWithSampler(
margin=0.5, sampler_inbatch=sampler_inbatch)
# 4. training with catalyst Runner
callbacks = [
dl.ControlFlowCallback(
dl.CriterionCallback(input_key="embeddings",
target_key="targets",
metric_key="loss"),
loaders="train",
),
dl.ControlFlowCallback(
dl.ReidCMCScoreCallback(
embeddings_key="embeddings",
pids_key="targets",
cids_key="cids",
is_query_key="is_query",
topk_args=[1],
),
loaders="valid",
),
dl.PeriodicLoaderCallback(valid_loader_key="valid",
valid_metric_key="cmc01",
minimize=False,
valid=2),
]
runner = ReIDCustomRunner()
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
callbacks=callbacks,
loaders=OrderedDict({
"train": train_loader,
"valid": valid_loader
}),
verbose=False,
logdir=logdir,
valid_loader="valid",
valid_metric="cmc01",
minimize_valid_metric=False,
num_epochs=10,
)
assert "cmc01" in runner.loader_metrics
assert runner.loader_metrics["cmc01"] > 0.7
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
# 1. train and valid loaders
transforms = Compose([ToTensor(), Normalize((0.1307, ), (0.3081, ))])
train_dataset = datasets.MnistMLDataset(root=os.getcwd(),
download=True,
transform=transforms)
sampler = data.BatchBalanceClassSampler(
labels=train_dataset.get_labels(),
num_classes=5,
num_samples=10,
num_batches=10)
train_loader = DataLoader(dataset=train_dataset, batch_sampler=sampler)
valid_dataset = datasets.MnistQGDataset(root=os.getcwd(),
transform=transforms,
gallery_fraq=0.2)
valid_loader = DataLoader(dataset=valid_dataset, batch_size=1024)
# 2. model and optimizer
model = models.MnistSimpleNet(out_features=16)
optimizer = Adam(model.parameters(), lr=0.001)
# 3. criterion with triplets sampling
sampler_inbatch = data.HardTripletsSampler(norm_required=False)
criterion = nn.TripletMarginLossWithSampler(
margin=0.5, sampler_inbatch=sampler_inbatch)
# 4. training with catalyst Runner
callbacks = [
dl.ControlFlowCallback(
dl.CriterionCallback(input_key="embeddings",
target_key="targets",
metric_key="loss"),
loaders="train",
),
dl.ControlFlowCallback(
dl.CMCScoreCallback(
embeddings_key="embeddings",
labels_key="targets",
is_query_key="is_query",
topk_args=[1],
),
loaders="valid",
),
dl.PeriodicLoaderCallback(valid_loader_key="valid",
valid_metric_key="cmc01",
minimize=False,
valid=2),
]
runner = CustomRunner(input_key="features", output_key="embeddings")
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
callbacks=callbacks,
loaders={
"train": train_loader,
"valid": valid_loader
},
verbose=False,
logdir=logdir,
valid_loader="valid",
valid_metric="cmc01",
minimize_valid_metric=False,
num_epochs=2,
)
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
# 1. data and transforms
transforms = Compose([
torchvision.transforms.ToPILImage(),
torchvision.transforms.RandomCrop((28, 28)),
torchvision.transforms.RandomVerticalFlip(),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
Normalize((0.1307, ), (0.3081, )),
])
transform_original = Compose([
ToTensor(),
Normalize((0.1307, ), (0.3081, )),
])
mnist = MNIST("./logdir", train=True, download=True, transform=None)
contrastive_mnist = SelfSupervisedDatasetWrapper(
mnist,
transforms=transforms,
transform_original=transform_original)
train_loader = torch.utils.data.DataLoader(contrastive_mnist,
batch_size=BATCH_SIZE)
mnist_valid = MNIST("./logdir",
train=False,
download=True,
transform=None)
contrastive_valid = SelfSupervisedDatasetWrapper(
mnist_valid,
transforms=transforms,
transform_original=transform_original)
valid_loader = torch.utils.data.DataLoader(contrastive_valid,
batch_size=BATCH_SIZE)
# 2. model and optimizer
encoder = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 16),
nn.LeakyReLU(inplace=True))
projection_head = nn.Sequential(
nn.Linear(16, 16, bias=False),
nn.ReLU(inplace=True),
nn.Linear(16, 16, bias=True),
)
class ContrastiveModel(torch.nn.Module):
def __init__(self, model, encoder):
super(ContrastiveModel, self).__init__()
self.model = model
self.encoder = encoder
def forward(self, x):
emb = self.encoder(x)
projection = self.model(emb)
return emb, projection
model = ContrastiveModel(model=projection_head, encoder=encoder)
optimizer = Adam(model.parameters(), lr=LR)
# 3. criterion with triplets sampling
criterion = NTXentLoss(tau=0.1)
callbacks = [
dl.ControlFlowCallback(
dl.CriterionCallback(input_key="projection_left",
target_key="projection_right",
metric_key="loss"),
loaders="train",
),
dl.SklearnModelCallback(
feature_key="embedding_left",
target_key="target",
train_loader="train",
valid_loaders="valid",
model_fn=RandomForestClassifier,
predict_method="predict_proba",
predict_key="sklearn_predict",
random_state=RANDOM_STATE,
n_estimators=50,
),
dl.ControlFlowCallback(
dl.AccuracyCallback(target_key="target",
input_key="sklearn_predict",
topk_args=(1, 3)),
loaders="valid",
),
]
runner = dl.SelfSupervisedRunner()
logdir = "./logdir"
runner.train(
model=model,
engine=engine or dl.DeviceEngine(device),
criterion=criterion,
optimizer=optimizer,
callbacks=callbacks,
loaders={
"train": train_loader,
"valid": valid_loader
},
verbose=False,
logdir=logdir,
valid_loader="train",
valid_metric="loss",
minimize_valid_metric=True,
num_epochs=TRAIN_EPOCH,
)
valid_path = Path(logdir) / "logs/valid.csv"
best_accuracy = max(
float(row["accuracy"]) for row in read_csv(valid_path)
if row["accuracy"] != "accuracy")
assert best_accuracy > 0.6
def forward(self, x):
emb = self.encoder(x)
projection = self.model(emb)
return emb, projection
model = ContrastiveModel(projection_head, encoder)
# 2. model and optimizer
optimizer = Adam(model.parameters(), lr=args.learning_rate)
# 3. criterion with triplets sampling
criterion = NTXentLoss(tau=args.temperature)
callbacks = [
dl.ControlFlowCallback(
dl.CriterionCallback(input_key="projection_left",
target_key="projection_right",
metric_key="loss"),
loaders="train",
)
]
runner = SelfSupervisedRunner()
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
callbacks=callbacks,
loaders={
"train": train_loader,
# "valid": valid_loader
},
