0de148f6892361e61d4b68cf97814e73692f6b1a,docs/_static/examples/gan.py,,,#,124

Before Change

// Configure data loader
os.makedirs("./data/mnist", exist_ok=True)
transform = transforms.Compose([
                        transforms.ToTensor(),
                        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                   ])

dataset = datasets.MNIST("./data/mnist", train=True, download=True, transform=transform)

dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)


// Model and optimizer

After Change

from torchbearer.metrics import mean, running_mean
metrics = ["loss", mean(running_mean(D_LOSS)), mean(running_mean(G_LOSS))]

trial = tb.Trial(generator, None, criterion=gen_crit, metrics=metrics, callbacks=[saver_callback])
trial.with_train_generator(dataloader, steps=200000)
trial.to(device)

new_keys = {DISC_MODEL: discriminator.to(device), DISC_OPT: optimizer_D, GEN_OPT: optimizer_G, DISC_CRIT: disc_crit}
trial.state.update(new_keys)
trial.with_closure(closure)
trial.run(epochs=1)

Instances