In [ ]:

%reload_ext autoreload
%autoreload 2

In [ ]:

from fastai import *
from fastai.vision import *

LSun bedroom data¶

For this lesson, we'll be using the bedrooms from the LSUN dataset. The full dataset is a bit too large so we'll use a sample from kaggle.

In [ ]:

path = untar_data(URLs.LSUN_BEDROOMS)

We then grab all the images in the folder with the data block API. We don't create a validation set here for reasons we'll explain later.

In [ ]:

class NoisyItem(ItemBase):
    def __init__(self, noise_sz): self.obj,self.data = noise_sz,torch.randn(noise_sz, 1, 1)
    def __str__(self):  return ''
    def apply_tfms(self, tfms, **kwargs): return self

In [ ]:

class GANItemList(ImageItemList):
    _label_cls = ImageItemList
    
    def __init__(self, items, noise_sz:int=100, **kwargs):
        super().__init__(items, **kwargs)
        self.noise_sz = noise_sz
        self.copy_new.append('noise_sz')
    
    def get(self, i): return NoisyItem(self.noise_sz)
    def reconstruct(self, t): return NoisyItem(t.size(0))
    
    def show_xys(self, xs, ys, imgsize:int=4, figsize:Optional[Tuple[int,int]]=None, **kwargs):
        super().show_xys(ys, xs, imgsize=imgsize, figsize=figsize, **kwargs)

In [ ]:

def get_data(bs, size):
    train_ds = (GANItemList.from_folder(path).label_from_func(noop)
               .transform(tfms=[crop_pad(size=size, row_pct=(0,1), col_pct=(0,1))], size=size, tfm_y=True))
    return (ImageDataBunch.create(train_ds, valid_ds=None, path=path, bs=bs)
                     .normalize(do_x=False, stats = [torch.tensor([0.5,0.5,0.5]), torch.tensor([0.5,0.5,0.5])], do_y=True))

We'll begin with a small side and use gradual resizing.

In [ ]:

data = get_data(128, 64)

In [ ]:

data.show_batch(rows=5)

Models¶

GAN stands for Generative Adversarial Nets and were invented by Ian Goodfellow. The concept is that we will train two models at the same time: a generator and a discriminator. The generator will try to make new images similar to the ones in our dataset, and the discriminator job will try to classify real images from the ones the generator does. The generator returns images, the discriminator a single number (usually 0. for fake images and 1. for real ones).

We train them against each other in the sense that at each step (more or less), we:

Freeze the generator and train the discriminator for one step by:

getting one batch of true images (let's call that real)
generating one batch of fake images (let's call that fake)
have the discriminator evaluate each batch and compute a loss function from that; the important part is that it rewards positively the detection of real images and penalizes the fake ones
update the weights of the discriminator with the gradients of this loss

Freeze the discriminator and train the generator for one step by:

generating one batch of fake images
evaluate the discriminator on it
return a loss that rewards posisitivly the discriminator thinking those are real images; the important part is that it rewards positively the detection of real images and penalizes the fake ones
update the weights of the generator with the gradients of this loss

Here, we'll use the Wassertein GAN.

We create a generator and a discriminator that we pass to gan_learner. The noise_size is the size of the random vector from which our generator creates images.

In [ ]:

from fastai.vision.gan import *

In [ ]:

generator = basic_generator(in_size=64, n_channels=3, n_extra_layers=1)
critic = basic_critic(in_size=64, n_channels=3, n_extra_layers=1)

In [ ]:

learn = GANLearner.wgan(data, generator, critic, opt_func=optim.RMSprop, wd=0.)

In [ ]:

learn.fit(1,1e-4)

In [ ]: