def test_pruning_callback() -> None:
"""Quantize model"""
loaders = {
"train": DataLoader(
MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()), batch_size=32
),
"valid": DataLoader(
MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()), batch_size=32
),
}
model = nn.Sequential(Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10))
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
runner = dl.SupervisedRunner()
runner.train(
model=model,
callbacks=[dl.QuantizationCallback(logdir="./logs")],
loaders=loaders,
criterion=criterion,
optimizer=optimizer,
num_epochs=1,
logdir="./logs",
check=True,
)
assert os.path.isfile("./logs/quantized.pth")
def test_balance_class_sampler_with_prefetch():
train_data = MNIST(os.getcwd(),
train=True,
download=True,
transform=ToTensor())
train_labels = train_data.targets.cpu().numpy().tolist()
train_sampler = BalanceClassSampler(train_labels, mode=5000)
valid_data = MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor())
loaders = {
"train": DataLoader(train_data, sampler=train_sampler, batch_size=32),
"valid": DataLoader(valid_data, batch_size=32),
}
loaders = {k: BatchPrefetchLoaderWrapper(v) for k, v in loaders.items()}
model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10))
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.02)
runner = dl.SupervisedRunner()
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
loaders=loaders,
num_epochs=1,
logdir="./logs",
valid_loader="valid",
valid_metric="loss",
minimize_valid_metric=True,
verbose=True,
)
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10))
optimizer = optim.Adam(model.parameters(), lr=0.02)
loaders = {
"train":
DataLoader(MNIST(os.getcwd(),
train=True,
download=True,
transform=ToTensor()),
batch_size=32),
"valid":
DataLoader(MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor()),
batch_size=32),
}
runner = CustomRunner()
# model training
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
optimizer=optimizer,
loaders=loaders,
logdir=logdir,
num_epochs=1,
verbose=False,
valid_loader="valid",
valid_metric="loss",
minimize_valid_metric=True,
)
def test_accuracy():
"""Test if accuracy drops too low."""
model = torch.nn.Sequential(
Flatten(),
torch.nn.Linear(28 * 28, 128),
torch.nn.ReLU(),
torch.nn.Linear(128, 64),
torch.nn.Linear(64, 10),
)
datasets = {
"train": MNIST("./data", transform=ToTensor(), download=True),
"valid": MNIST("./data", transform=ToTensor(), train=False),
}
dataloaders = {
k: torch.utils.data.DataLoader(d, batch_size=32)
for k, d in datasets.items()
}
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
runner = SupervisedRunner()
runner.train(
model=model,
optimizer=optimizer,
loaders=dataloaders,
callbacks=[AccuracyCallback(target_key="targets", input_key="logits")],
num_epochs=1,
criterion=torch.nn.CrossEntropyLoss(),
valid_loader="valid",
valid_metric="accuracy01",
minimize_valid_metric=False,
)
accuracy_before = _evaluate_loader_accuracy(runner, dataloaders["valid"])
q_model = quantize_model(model)
runner.model = q_model
accuracy_after = _evaluate_loader_accuracy(runner, dataloaders["valid"])
assert abs(accuracy_before - accuracy_after) < 0.01
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
teacher = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10))
student = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10))
model = {"teacher": teacher, "student": student}
criterion = {"cls": nn.CrossEntropyLoss(), "kl": nn.KLDivLoss(reduction="batchmean")}
optimizer = optim.Adam(student.parameters(), lr=0.02)
loaders = {
"train": DataLoader(
MNIST(os.getcwd(), train=True, download=True, transform=ToTensor()), batch_size=32
),
"valid": DataLoader(
MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()), batch_size=32
),
}
runner = DistilRunner()
# model training
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
loaders=loaders,
num_epochs=1,
logdir=logdir,
verbose=False,
callbacks=[
dl.AccuracyCallback(
input_key="t_logits", target_key="targets", num_classes=2, prefix="teacher_"
),
dl.AccuracyCallback(
input_key="s_logits", target_key="targets", num_classes=2, prefix="student_"
),
dl.CriterionCallback(
input_key="s_logits",
target_key="targets",
metric_key="cls_loss",
criterion_key="cls",
),
dl.CriterionCallback(
input_key="s_logprobs",
target_key="t_probs",
metric_key="kl_div_loss",
criterion_key="kl",
),
dl.MetricAggregationCallback(
metric_key="loss", metrics=["kl_div_loss", "cls_loss"], mode="mean"
),
dl.OptimizerCallback(metric_key="loss", model_key="student"),
dl.CheckpointCallback(
logdir=logdir,
loader_key="valid",
metric_key="loss",
minimize=True,
save_n_best=3,
),
],
)
def get_loaders(self, stage: str):
loaders = {
"train": DataLoader(
MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()), batch_size=32
),
"valid": DataLoader(
MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()), batch_size=32
),
}
return loaders
def test_transform_kornia():
"""Run few epochs to check ``BatchTransformCallback`` callback."""
model = torch.nn.Linear(28 * 28, 10)
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
loaders = {
"train":
DataLoader(
MnistDataset(
MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor())),
batch_size=32,
),
"valid":
DataLoader(
MnistDataset(
MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor())),
batch_size=32,
),
}
transrorms = [
augmentation.RandomAffine(degrees=(-15, 20), scale=(0.75, 1.25)),
]
runner = CustomRunner()
# model training
runner.train(
model=model,
optimizer=optimizer,
loaders=loaders,
logdir="./logs",
num_epochs=5,
verbose=False,
load_best_on_end=True,
check=True,
callbacks=[
BatchTransformCallback(transform=transrorms,
scope="on_batch_start",
input_key="features")
],
)
# model inference
for prediction in runner.predict_loader(loader=loaders["train"]):
assert prediction.detach().cpu().numpy().shape[-1] == 10
def objective(trial):
lr = trial.suggest_loguniform("lr", 1e-3, 1e-1)
num_hidden = int(trial.suggest_loguniform("num_hidden", 32, 128))
loaders = {
"train":
DataLoader(
MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor()),
batch_size=32,
),
"valid":
DataLoader(
MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor()),
batch_size=32,
),
}
model = nn.Sequential(nn.Flatten(), nn.Linear(784, num_hidden),
nn.ReLU(), nn.Linear(num_hidden, 10))
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
runner = dl.SupervisedRunner(input_key="features",
output_key="logits",
target_key="targets")
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
loaders=loaders,
callbacks={
"optuna":
dl.OptunaPruningCallback(loader_key="valid",
metric_key="accuracy01",
minimize=False,
trial=trial),
"accuracy":
dl.AccuracyCallback(input_key="logits",
target_key="targets",
num_classes=10),
},
num_epochs=2,
)
score = trial.best_score
return score
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
# <--- multi-model setup --->
encoder = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 128))
head = nn.Linear(128, 10)
model = {"encoder": encoder, "head": head}
optimizer = optim.Adam([
{
"params": encoder.parameters()
},
{
"params": head.parameters()
},
],
lr=0.02)
# <--- multi-model setup --->
criterion = nn.CrossEntropyLoss()
loaders = {
"train":
DataLoader(MNIST(os.getcwd(),
train=True,
download=True,
transform=ToTensor()),
batch_size=32),
"valid":
DataLoader(MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor()),
batch_size=32),
}
runner = CustomRunner()
# model training
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
loaders=loaders,
logdir=logdir,
num_epochs=1,
verbose=False,
valid_loader="valid",
valid_metric="loss",
minimize_valid_metric=True,
)
def test_optuna():
trainset = MNIST("./data",
train=False,
download=True,
transform=ToTensor())
testset = MNIST("./data", train=False, download=True, transform=ToTensor())
loaders = {
"train": DataLoader(trainset, batch_size=32),
"valid": DataLoader(testset, batch_size=64),
}
model = nn.Sequential(Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))
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
study = optuna.create_study(
direction="maximize",
pruner=optuna.pruners.MedianPruner(n_startup_trials=1,
n_warmup_steps=0,
interval_steps=1),
)
study.optimize(objective, n_trials=2, timeout=300)
assert True
from catalyst.settings import IS_CUDA_AVAILABLE, NUM_CUDA_DEVICES, SETTINGS
if SETTINGS.cv_required:
import torchvision
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, )),
])
class ProjectionHead(nn.Module):
def __init__(self):
super(ProjectionHead, self).__init__()
self.seq = nn.Sequential(
nn.Linear(16, 16, bias=False),
nn.ReLU(inplace=True),
nn.Linear(16, 16, bias=True),
)
def forward(self, x):
return self.seq(x)
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
# latent_dim = 128
# generator = nn.Sequential(
# # We want to generate 128 coefficients to reshape into a 7x7x128 map
# nn.Linear(128, 128 * 7 * 7),
# nn.LeakyReLU(0.2, inplace=True),
# Lambda(lambda x: x.view(x.size(0), 128, 7, 7)),
# nn.ConvTranspose2d(128, 128, (4, 4), stride=(2, 2), padding=1),
# nn.LeakyReLU(0.2, inplace=True),
# nn.ConvTranspose2d(128, 128, (4, 4), stride=(2, 2), padding=1),
# nn.LeakyReLU(0.2, inplace=True),
# nn.Conv2d(128, 1, (7, 7), padding=3),
# nn.Sigmoid(),
# )
# discriminator = nn.Sequential(
# nn.Conv2d(1, 64, (3, 3), stride=(2, 2), padding=1),
# nn.LeakyReLU(0.2, inplace=True),
# nn.Conv2d(64, 128, (3, 3), stride=(2, 2), padding=1),
# nn.LeakyReLU(0.2, inplace=True),
# GlobalMaxPool2d(),
# Flatten(),
# nn.Linear(128, 1),
# )
latent_dim = 32
generator = nn.Sequential(
nn.Linear(latent_dim, 28 * 28),
Lambda(_ddp_hack),
nn.Sigmoid(),
)
discriminator = nn.Sequential(Flatten(), nn.Linear(28 * 28, 1))
model = {"generator": generator, "discriminator": discriminator}
criterion = {
"generator": nn.BCEWithLogitsLoss(),
"discriminator": nn.BCEWithLogitsLoss()
}
optimizer = {
"generator":
torch.optim.Adam(generator.parameters(),
lr=0.0003,
betas=(0.5, 0.999)),
"discriminator":
torch.optim.Adam(discriminator.parameters(),
lr=0.0003,
betas=(0.5, 0.999)),
}
loaders = {
"train":
DataLoader(MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor()),
batch_size=32),
}
runner = CustomRunner(latent_dim)
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
loaders=loaders,
callbacks=[
dl.CriterionCallback(
input_key="combined_predictions",
target_key="labels",
metric_key="loss_discriminator",
criterion_key="discriminator",
),
dl.CriterionCallback(
input_key="generated_predictions",
target_key="misleading_labels",
metric_key="loss_generator",
criterion_key="generator",
),
dl.OptimizerCallback(
model_key="generator",
optimizer_key="generator",
metric_key="loss_generator",
),
dl.OptimizerCallback(
model_key="discriminator",
optimizer_key="discriminator",
metric_key="loss_discriminator",
),
],
valid_loader="train",
valid_metric="loss_generator",
minimize_valid_metric=True,
num_epochs=1,
verbose=False,
logdir=logdir,
)
if not isinstance(engine, dl.DistributedDataParallelEngine):
runner.predict_batch(None)[0, 0].cpu().numpy()
def get_transform(self, stage: str = None, mode: str = None):
return ToTensor()
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
from catalyst.settings import IS_CUDA_AVAILABLE, NUM_CUDA_DEVICES, SETTINGS
if SETTINGS.cv_required:
import torchvision
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, ))])
class ProjectionHead(nn.Module):
def __init__(self):
super(ProjectionHead, self).__init__()
self.seq = nn.Sequential(nn.Linear(16, 16, bias=False),
nn.ReLU(inplace=True),
nn.Linear(16, 16, bias=True))
def forward(self, x):
return self.seq(x)
class Encoder(nn.Module):
def __init__(self):
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10))
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.02)
loaders = {
"train":
DataLoader(MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor()),
batch_size=32),
"valid":
DataLoader(MNIST(os.getcwd(),
train=False,
download=True,
transform=ToTensor()),
batch_size=32),
}
runner = dl.SupervisedRunner(input_key="features",
output_key="logits",
target_key="targets",
loss_key="loss")
callbacks = [
dl.AccuracyCallback(input_key="logits",
target_key="targets",
topk_args=(1, 3, 5)),
dl.PrecisionRecallF1SupportCallback(input_key="logits",
target_key="targets",
num_classes=10),
]
if SETTINGS.ml_required:
callbacks.append(
dl.ConfusionMatrixCallback(input_key="logits",
target_key="targets",
num_classes=10))
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"))
if SETTINGS.onnx_required:
callbacks.append(
dl.OnnxCallback(logdir=logdir, input_key="features"))
if SETTINGS.pruning_required:
callbacks.append(
dl.PruningCallback(pruning_fn="l1_unstructured", amount=0.5))
if SETTINGS.quantization_required:
callbacks.append(dl.QuantizationCallback(logdir=logdir))
if engine is None or not isinstance(engine,
dl.DistributedDataParallelEngine):
callbacks.append(
dl.TracingCallback(logdir=logdir, input_key="features"))
# model training
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
loaders=loaders,
num_epochs=1,
callbacks=callbacks,
logdir=logdir,
valid_loader="valid",
valid_metric="loss",
minimize_valid_metric=True,
verbose=False,
load_best_on_end=True,
timeit=False,
check=False,
overfit=False,
fp16=False,
ddp=False,
)
# loader evaluation
metrics = runner.evaluate_loader(
model=runner.model,
loader=loaders["valid"],
callbacks=[
dl.AccuracyCallback(input_key="logits",
target_key="targets",
topk_args=(1, 3, 5))
],
)
assert "accuracy" in metrics.keys()
# model inference
for prediction in runner.predict_loader(loader=loaders["valid"]):
assert prediction["logits"].detach().cpu().numpy().shape[-1] == 10
# model post-processing
features_batch = next(iter(loaders["valid"]))[0]
# model stochastic weight averaging
model.load_state_dict(
utils.get_averaged_weights_by_path_mask(logdir=logdir,
path_mask="*.pth"))
# model onnx export
if SETTINGS.onnx_required:
utils.onnx_export(
model=runner.model,
batch=runner.engine.sync_device(features_batch),
file="./mnist.onnx",
verbose=False,
)
# model quantization
if SETTINGS.quantization_required:
utils.quantize_model(model=runner.model)
# model pruning
if SETTINGS.pruning_required:
utils.prune_model(model=runner.model,
pruning_fn="l1_unstructured",
amount=0.8)
# model tracing
utils.trace_model(model=runner.model, batch=features_batch)