def test_pruning_callback() -> None:
"""Quantize model"""
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,
),
}
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.DynamicQuantizationCallback(out_dir="./logs", backend="qnnpack")
],
loaders=loaders,
criterion=criterion,
optimizer=optimizer,
num_epochs=1,
logdir="./logs",
check=True,
)
assert os.path.isfile("./logs/best_quantized.pth")
def datasets_fn():
"""
Docs.
"""
dataset = MNIST(
"./data", train=False, download=True, transform=ToTensor(),
)
return {"train": dataset, "valid": dataset}
def simple_transform():
"""
Simple Transform
Returns:
Compose: transforms
"""
return Compose([ToTensor(), Normalize((0.1307, ), (0.3081, ))])
def _get_loaders(*, root: str, batch_size: int = 1, num_workers: int = 1):
data_transform = ToTensor()
trainset = MNIST(root=root,
train=True,
download=True,
transform=data_transform)
trainloader = DataLoader(trainset,
batch_size=batch_size,
num_workers=num_workers)
testset = MNIST(root=root,
train=False,
download=True,
transform=data_transform)
testloader = DataLoader(testset,
batch_size=batch_size,
num_workers=num_workers)
loaders = collections.OrderedDict(train=trainloader, valid=testloader)
return loaders
def main():
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),
)
model = {"generator": generator, "discriminator": discriminator}
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(
"./data",
train=True,
download=True,
transform=ToTensor(),
),
batch_size=32,
),
}
runner = CustomRunner()
runner.train(
model=model,
optimizer=optimizer,
loaders=loaders,
callbacks=[
dl.OptimizerCallback(optimizer_key="generator",
metric_key="loss_generator"),
dl.OptimizerCallback(optimizer_key="discriminator",
metric_key="loss_discriminator"),
],
main_metric="loss_generator",
num_epochs=20,
verbose=True,
logdir="./logs_gan",
check=True,
)
nn.Linear(128, 1))
model = {"generator": generator, "discriminator": discriminator}
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=True,
download=True,
transform=ToTensor()),
batch_size=32),
}
class CustomRunner(dl.Runner):
def _handle_batch(self, batch):
real_images, _ = batch
batch_metrics = {}
# Sample random points in the latent space
batch_size = real_images.shape[0]
random_latent_vectors = torch.randn(batch_size,
latent_dim).to(self.device)
# Decode them to fake images
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from catalyst import dl, metrics
from catalyst.contrib.data.cv import ToTensor
from catalyst.contrib.datasets import MNIST
model = torch.nn.Linear(28 * 28, 10)
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
loaders = {
"train":
DataLoader(
MNIST("./data", train=True, download=True, transform=ToTensor()),
batch_size=32,
),
}
class CustomRunner(dl.Runner):
def predict_batch(self, batch):
# model inference step
return self.model(batch[0].to(self.device).view(batch[0].size(0), -1))
def _handle_batch(self, batch):
# model train/valid step
x, y = batch
y_hat = self.model(x.view(x.size(0), -1))