In [1]:
%reload_ext autoreload
%autoreload 2
In [2]:
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 [3]:
path = Config.data_path()/'lsun'
path.mkdir(parents=True, exist_ok=True)
path
Out[3]:
PosixPath('/home/ubuntu/.fastai/data/lsun')

Uncomment the next commands to download and extract the data in your machine.

In [4]:
#! kaggle datasets download -d jhoward/lsun_bedroom -p {path}  
In [5]:
#! unzip -q -n {path}/lsun_bedroom.zip -d {path}
#! unzip -q -n {path}/sample.zip -d {path}

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 [6]:
def get_data(bs, size):
    train_ds = (ImageItemList.from_folder(path).label_empty()
               .transform(tfms=[crop_pad(size=size, row_pct=(0,1), col_pct=(0,1))],size=size))
    return (ImageDataBunch.create(train_ds, valid_ds=None, path=path, bs=bs)
                     .normalize(stats = [torch.tensor([0.5,0.5,0.5]), torch.tensor([0.5,0.5,0.5])]))

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

In [7]:
data = get_data(64, 64)
In [8]:
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:

  1. 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
  1. 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.

In [9]:
def gan_learner(data, generator, discriminator, loss_funcD=None, loss_funcG=None, noise_size:int=None, wgan:bool=False,
                **kwargs):
    gan = models.GAN(generator, discriminator)
    learn = Learner(data, gan, loss_func=NoopLoss(), **kwargs)
    if wgan: loss_funcD,loss_funcG = WasserteinLoss(),noop
    if noise_size is None: cb = callbacks.GANTrainer(learn, loss_funcD, loss_funcG)
    else: cb = callbacks.NoisyGANTrainer(learn, loss_funcD, loss_funcG, bs=data.batch_size, noise_sz=noise_size)
    learn.callbacks.append(cb)
    return learn
In [10]:
generator = models.basic_generator(in_size=64, n_channels=3, n_extra_layers=1)
discriminator = models.basic_discriminator(in_size=64, n_channels=3, n_extra_layers=1)
In [11]:
learn = gan_learner(data, generator, discriminator, wgan=True, noise_size=100,
                    opt_func=optim.RMSprop, wd=0.)
In [12]:
learn.fit(1, 1e-4)
Total time: 15:05
epoch  train_loss  valid_loss  gen_loss  disc_loss
1      -0.716458               0.670430  -1.410371  (15:05)

In [13]:
x,y = next(iter(learn.data.train_dl))
In [20]:
tst = learn.model(learn.callbacks[0].input_fake(x, grad=False), gen=True)
img = data.denorm(tst[0].cpu()).numpy().clip(0,1).transpose(1,2,0)
plt.imshow(img)
Out[20]:
<matplotlib.image.AxesImage at 0x7fb6dc85b048>
In [14]:
learn.model(tst)
Out[14]:
tensor([-0.1581], device='cuda:0', grad_fn=<ViewBackward>)
In [15]:
learn.model(x)
Out[15]:
tensor([-0.2798], device='cuda:0', grad_fn=<ViewBackward>)
In [ ]: