chore: import upstream snapshot with attribution

python/ray/tune/examples/pbt_dcgan_mnist/common.py

@@ -0,0 +1,285 @@
+import os
+
+import matplotlib.animation as animation
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.utils.data
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+import torchvision.utils as vutils
+from scipy.stats import entropy
+from torch.autograd import Variable
+from torch.nn import functional as F
+
+import ray
+
+# Training parameters
+workers = 2
+batch_size = 64
+image_size = 32
+
+# Number of channels in the training images. For color images this is 3
+nc = 1
+
+# Size of z latent vector (i.e. size of generator input)
+nz = 100
+
+# Size of feature maps in generator
+ngf = 32
+
+# Size of feature maps in discriminator
+ndf = 32
+
+# Beta1 hyperparam for Adam optimizers
+beta1 = 0.5
+
+# iterations of actual training in each Trainable _train
+train_iterations_per_step = 5
+
+MODEL_PATH = os.path.expanduser("~/.ray/models/mnist_cnn.pt")
+
+
+def get_data_loader(data_dir="~/data"):
+    dataset = dset.MNIST(
+        root=data_dir,
+        download=True,
+        transform=transforms.Compose(
+            [
+                transforms.Resize(image_size),
+                transforms.ToTensor(),
+                transforms.Normalize((0.13066,), (0.30131,)),
+            ]
+        ),
+    )
+
+    # Create the dataloader
+    dataloader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=True, num_workers=workers
+    )
+
+    return dataloader
+
+
+# __GANmodel_begin__
+# custom weights initialization called on netG and netD
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find("BatchNorm") != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+
+# Generator Code
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.main = nn.Sequential(
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(nz, ngf * 4, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(ngf * 4),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf * 2),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
+            nn.Tanh(),
+        )
+
+    def forward(self, input):
+        return self.main(input)
+
+
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+        self.main = nn.Sequential(
+            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 2),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 4),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, input):
+        return self.main(input)
+
+
+# __GANmodel_end__
+
+
+# __INCEPTION_SCORE_begin__
+class Net(nn.Module):
+    """
+    LeNet for MNist classification, used for inception_score
+    """
+
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def inception_score(imgs, mnist_model_ref, batch_size=32, splits=1):
+    N = len(imgs)
+    dtype = torch.FloatTensor
+    dataloader = torch.utils.data.DataLoader(imgs, batch_size=batch_size)
+    cm = ray.get(mnist_model_ref)  # Get the mnist model from Ray object store.
+    up = nn.Upsample(size=(28, 28), mode="bilinear").type(dtype)
+
+    def get_pred(x):
+        x = up(x)
+        x = cm(x)
+        return F.softmax(x).data.cpu().numpy()
+
+    preds = np.zeros((N, 10))
+    for i, batch in enumerate(dataloader, 0):
+        batch = batch.type(dtype)
+        batchv = Variable(batch)
+        batch_size_i = batch.size()[0]
+        preds[i * batch_size : i * batch_size + batch_size_i] = get_pred(batchv)
+
+    # Now compute the mean kl-div
+    split_scores = []
+    for k in range(splits):
+        part = preds[k * (N // splits) : (k + 1) * (N // splits), :]
+        py = np.mean(part, axis=0)
+        scores = []
+        for i in range(part.shape[0]):
+            pyx = part[i, :]
+            scores.append(entropy(pyx, py))
+        split_scores.append(np.exp(np.mean(scores)))
+
+    return np.mean(split_scores), np.std(split_scores)
+
+
+# __INCEPTION_SCORE_end__
+
+
+def train_func(
+    netD,
+    netG,
+    optimG,
+    optimD,
+    criterion,
+    dataloader,
+    iteration,
+    device,
+    mnist_model_ref,
+):
+    real_label = 1
+    fake_label = 0
+
+    for i, data in enumerate(dataloader, 0):
+        if i >= train_iterations_per_step:
+            break
+
+        netD.zero_grad()
+        real_cpu = data[0].to(device)
+        b_size = real_cpu.size(0)
+        label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
+        output = netD(real_cpu).view(-1)
+        errD_real = criterion(output, label)
+        errD_real.backward()
+        D_x = output.mean().item()
+
+        noise = torch.randn(b_size, nz, 1, 1, device=device)
+        fake = netG(noise)
+        label.fill_(fake_label)
+        output = netD(fake.detach()).view(-1)
+        errD_fake = criterion(output, label)
+        errD_fake.backward()
+        D_G_z1 = output.mean().item()
+        errD = errD_real + errD_fake
+        optimD.step()
+
+        netG.zero_grad()
+        label.fill_(real_label)
+        output = netD(fake).view(-1)
+        errG = criterion(output, label)
+        errG.backward()
+        D_G_z2 = output.mean().item()
+        optimG.step()
+
+        is_score, is_std = inception_score(fake, mnist_model_ref)
+
+        # Output training stats
+        if iteration % 10 == 0:
+            print(
+                "[%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z))"
+                ": %.4f / %.4f \tInception score: %.4f"
+                % (
+                    iteration,
+                    len(dataloader),
+                    errD.item(),
+                    errG.item(),
+                    D_x,
+                    D_G_z1,
+                    D_G_z2,
+                    is_score,
+                )
+            )
+
+    return errG.item(), errD.item(), is_score
+
+
+def plot_images(dataloader):
+    # Plot some training images
+    real_batch = next(iter(dataloader))
+    plt.figure(figsize=(8, 8))
+    plt.axis("off")
+    plt.title("Original Images")
+    plt.imshow(
+        np.transpose(
+            vutils.make_grid(real_batch[0][:64], padding=2, normalize=True).cpu(),
+            (1, 2, 0),
+        )
+    )
+
+    plt.show()
+
+
+def demo_gan(checkpoint_paths):
+    img_list = []
+    fixed_noise = torch.randn(64, nz, 1, 1)
+    for path in checkpoint_paths:
+        checkpoint_dict = torch.load(os.path.join(path, "checkpoint.pt"))
+
+        loadedG = Generator()
+        loadedG.load_state_dict(checkpoint_dict["netGmodel"])
+        with torch.no_grad():
+            fake = loadedG(fixed_noise).detach().cpu()
+        img_list.append(vutils.make_grid(fake, padding=2, normalize=True))
+
+    fig = plt.figure(figsize=(8, 8))
+    plt.axis("off")
+    ims = [[plt.imshow(np.transpose(i, (1, 2, 0)), animated=True)] for i in img_list]
+    ani = animation.ArtistAnimation(
+        fig, ims, interval=1000, repeat_delay=1000, blit=True
+    )
+    ani.save("./generated.gif", writer="imagemagick", dpi=72)
+    plt.show()

python/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_func.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python
+"""
+Example of training DCGAN on MNIST using PBT with Tune's function API.
+"""
+import argparse
+import os
+import tempfile
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.optim as optim
+import torch.utils.data
+from filelock import FileLock
+
+import ray
+from ray import tune
+from ray.tune import Checkpoint
+from ray.tune.examples.pbt_dcgan_mnist.common import (
+    MODEL_PATH,
+    Discriminator,
+    Generator,
+    Net,
+    beta1,
+    demo_gan,
+    get_data_loader,
+    plot_images,
+    train_func,
+    weights_init,
+)
+from ray.tune.schedulers import PopulationBasedTraining
+
+
+# __Train_begin__
+def dcgan_train(config):
+    use_cuda = config.get("use_gpu") and torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+    netD = Discriminator().to(device)
+    netD.apply(weights_init)
+    netG = Generator().to(device)
+    netG.apply(weights_init)
+    criterion = nn.BCELoss()
+    optimizerD = optim.Adam(
+        netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+    )
+    optimizerG = optim.Adam(
+        netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+    )
+    with FileLock(os.path.expanduser("~/ray_results/.data.lock")):
+        dataloader = get_data_loader()
+
+    step = 1
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            checkpoint_dict = torch.load(os.path.join(checkpoint_dir, "checkpoint.pt"))
+        netD.load_state_dict(checkpoint_dict["netDmodel"])
+        netG.load_state_dict(checkpoint_dict["netGmodel"])
+        optimizerD.load_state_dict(checkpoint_dict["optimD"])
+        optimizerG.load_state_dict(checkpoint_dict["optimG"])
+        # Note: Make sure to increment the loaded step by 1 to get the
+        # current step.
+        last_step = checkpoint_dict["step"]
+        step = last_step + 1
+
+        # NOTE: It's important to set the optimizer learning rates
+        # again, since we want to explore the parameters passed in by PBT.
+        # Without this, we would continue using the exact same
+        # configuration as the trial whose checkpoint we are exploiting.
+        if "netD_lr" in config:
+            for param_group in optimizerD.param_groups:
+                param_group["lr"] = config["netD_lr"]
+        if "netG_lr" in config:
+            for param_group in optimizerG.param_groups:
+                param_group["lr"] = config["netG_lr"]
+
+    while True:
+        lossG, lossD, is_score = train_func(
+            netD,
+            netG,
+            optimizerG,
+            optimizerD,
+            criterion,
+            dataloader,
+            step,
+            device,
+            config["mnist_model_ref"],
+        )
+        metrics = {"lossg": lossG, "lossd": lossD, "is_score": is_score}
+
+        if step % config["checkpoint_interval"] == 0:
+            with tempfile.TemporaryDirectory() as tmpdir:
+                torch.save(
+                    {
+                        "netDmodel": netD.state_dict(),
+                        "netGmodel": netG.state_dict(),
+                        "optimD": optimizerD.state_dict(),
+                        "optimG": optimizerG.state_dict(),
+                        "step": step,
+                    },
+                    os.path.join(tmpdir, "checkpoint.pt"),
+                )
+                tune.report(metrics, checkpoint=Checkpoint.from_directory(tmpdir))
+        else:
+            tune.report(metrics)
+
+        step += 1
+
+
+# __Train_end__
+
+
+def download_mnist_cnn():
+    import urllib.request
+
+    # Download a pre-trained MNIST model for inception score calculation.
+    # This is a tiny model (<100kb).
+    if not os.path.exists(MODEL_PATH):
+        print("downloading model")
+        os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
+        urllib.request.urlretrieve(
+            "https://github.com/ray-project/ray/raw/master/python/ray/tune/"
+            "examples/pbt_dcgan_mnist/mnist_cnn.pt",
+            MODEL_PATH,
+        )
+    return MODEL_PATH
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    parser.add_argument(
+        "--data-dir", type=str, default="~/data/", help="Set the path of the dataset."
+    )
+    args, _ = parser.parse_known_args()
+    ray.init()
+
+    download_mnist_cnn()
+
+    dataloader = get_data_loader(args.data_dir)
+    if not args.smoke_test:
+        plot_images(dataloader)
+
+    # __tune_begin__
+
+    # load the pretrained mnist classification model for inception_score
+    mnist_cnn = Net()
+    mnist_cnn.load_state_dict(torch.load(MODEL_PATH))
+    mnist_cnn.eval()
+    # Put the model in Ray object store.
+    mnist_model_ref = ray.put(mnist_cnn)
+
+    scheduler = PopulationBasedTraining(
+        perturbation_interval=5,
+        hyperparam_mutations={
+            # distribution for resampling
+            "netG_lr": lambda: np.random.uniform(1e-2, 1e-5),
+            "netD_lr": lambda: np.random.uniform(1e-2, 1e-5),
+        },
+    )
+
+    tune_iter = 5 if args.smoke_test else 300
+    tuner = tune.Tuner(
+        dcgan_train,
+        run_config=tune.RunConfig(
+            name="pbt_dcgan_mnist",
+            stop={"training_iteration": tune_iter},
+            verbose=1,
+        ),
+        tune_config=tune.TuneConfig(
+            metric="is_score",
+            mode="max",
+            num_samples=8,
+            scheduler=scheduler,
+        ),
+        param_space={
+            "netG_lr": tune.choice([0.0001, 0.0002, 0.0005]),
+            "netD_lr": tune.choice([0.0001, 0.0002, 0.0005]),
+            "mnist_model_ref": mnist_model_ref,
+        },
+    )
+    results = tuner.fit()
+    # __tune_end__
+
+    # demo of the trained Generators
+    if not args.smoke_test:
+        checkpoint_paths = [result.checkpoint.to_directory() for result in results]
+        demo_gan(checkpoint_paths)

python/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_trainable.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python
+"""
+Example of training DCGAN on MNIST using PBT with Tune's Trainable Class
+API.
+"""
+import argparse
+import os
+import random
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.optim as optim
+import torch.utils.data
+from filelock import FileLock
+
+import ray
+from ray import tune
+from ray.tune.examples.pbt_dcgan_mnist.common import (
+    MODEL_PATH,
+    Discriminator,
+    Generator,
+    Net,
+    beta1,
+    demo_gan,
+    get_data_loader,
+    plot_images,
+    train_func,
+    weights_init,
+)
+from ray.tune.schedulers import PopulationBasedTraining
+
+
+# __Trainable_begin__
+class PytorchTrainable(tune.Trainable):
+    def setup(self, config):
+        use_cuda = config.get("use_gpu") and torch.cuda.is_available()
+        self.device = torch.device("cuda" if use_cuda else "cpu")
+        self.netD = Discriminator().to(self.device)
+        self.netD.apply(weights_init)
+        self.netG = Generator().to(self.device)
+        self.netG.apply(weights_init)
+        self.criterion = nn.BCELoss()
+        self.optimizerD = optim.Adam(
+            self.netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+        )
+        self.optimizerG = optim.Adam(
+            self.netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+        )
+        with FileLock(os.path.expanduser("~/.data.lock")):
+            self.dataloader = get_data_loader(config.get("data_dir", "~/data"))
+        self.mnist_model_ref = config["mnist_model_ref"]
+
+    def step(self):
+        lossG, lossD, is_score = train_func(
+            self.netD,
+            self.netG,
+            self.optimizerG,
+            self.optimizerD,
+            self.criterion,
+            self.dataloader,
+            self._iteration,
+            self.device,
+            self.mnist_model_ref,
+        )
+        return {"lossg": lossG, "lossd": lossD, "is_score": is_score}
+
+    def save_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint.pt")
+        torch.save(
+            {
+                "netDmodel": self.netD.state_dict(),
+                "netGmodel": self.netG.state_dict(),
+                "optimD": self.optimizerD.state_dict(),
+                "optimG": self.optimizerG.state_dict(),
+            },
+            path,
+        )
+
+        return checkpoint_dir
+
+    def load_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint.pt")
+        checkpoint = torch.load(path)
+        self.netD.load_state_dict(checkpoint["netDmodel"])
+        self.netG.load_state_dict(checkpoint["netGmodel"])
+        self.optimizerD.load_state_dict(checkpoint["optimD"])
+        self.optimizerG.load_state_dict(checkpoint["optimG"])
+
+    def reset_config(self, new_config):
+        if "netD_lr" in new_config:
+            for param_group in self.optimizerD.param_groups:
+                param_group["lr"] = new_config["netD_lr"]
+        if "netG_lr" in new_config:
+            for param_group in self.optimizerG.param_groups:
+                param_group["lr"] = new_config["netG_lr"]
+
+        self.config = new_config
+        return True
+
+
+# __Trainable_end__
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    parser.add_argument(
+        "--data-dir", type=str, default="~/data/", help="Set the path of the dataset."
+    )
+    args, _ = parser.parse_known_args()
+    ray.init()
+
+    import urllib.request
+
+    # Download a pre-trained MNIST model for inception score calculation.
+    # This is a tiny model (<100kb).
+    if not os.path.exists(MODEL_PATH):
+        print("downloading model")
+        os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
+        urllib.request.urlretrieve(
+            "https://github.com/ray-project/ray/raw/master/python/ray/tune/"
+            "examples/pbt_dcgan_mnist/mnist_cnn.pt",
+            MODEL_PATH,
+        )
+
+    dataloader = get_data_loader()
+    if not args.smoke_test:
+        plot_images(dataloader)
+
+    # load the pretrained mnist classification model for inception_score
+    mnist_cnn = Net()
+    mnist_cnn.load_state_dict(torch.load(MODEL_PATH))
+    mnist_cnn.eval()
+    mnist_model_ref = ray.put(mnist_cnn)
+
+    # __tune_begin__
+    scheduler = PopulationBasedTraining(
+        time_attr="training_iteration",
+        perturbation_interval=5,
+        hyperparam_mutations={
+            # distribution for resampling
+            "netG_lr": lambda: np.random.uniform(1e-2, 1e-5),
+            "netD_lr": lambda: np.random.uniform(1e-2, 1e-5),
+        },
+    )
+
+    tune_iter = 10 if args.smoke_test else 300
+    tuner = tune.Tuner(
+        PytorchTrainable,
+        run_config=tune.RunConfig(
+            name="pbt_dcgan_mnist",
+            stop={"training_iteration": tune_iter},
+            verbose=1,
+            checkpoint_config=tune.CheckpointConfig(checkpoint_at_end=True),
+        ),
+        tune_config=tune.TuneConfig(
+            metric="is_score",
+            mode="max",
+            num_samples=8,
+            scheduler=scheduler,
+            reuse_actors=True,
+        ),
+        param_space={
+            "netG_lr": tune.sample_from(
+                lambda spec: random.choice([0.0001, 0.0002, 0.0005])
+            ),
+            "netD_lr": tune.sample_from(
+                lambda spec: random.choice([0.0001, 0.0002, 0.0005])
+            ),
+            "mnist_model_ref": mnist_model_ref,
+            "data_dir": args.data_dir,
+        },
+    )
+    results = tuner.fit()
+
+    # export_formats=[ExportFormat.MODEL]
+    # __tune_end__
+
+    # demo of the trained Generators
+    if not args.smoke_test:
+        checkpoint_paths = [result.checkpoint.to_directory() for result in results]
+        demo_gan(checkpoint_paths)