Training GANs
This is a short introductory post on training GANs using fastai2. This is mostly done to introduce the concept so that future notebooks are easier to understand.
from nbdev import export2html
from fastai2.vision.gan import *
from fastai2.basics import *
from fastai2.vision.all import *
For a better understanding of GANs, and how fastai2 does GAN training, take a look at the source notebook here: https://github.com/fastai/fastai2/blob/master/nbs/24_vision.gan.ipynb
GANs³ are generational adversarial networks, or in this case, a network than generates images, and another network,the critic, that tells if it is real or fake. The two networks then compete to be better than the other, and that is how we get a network that can generate images.
In this section we will setup a dataloader in fastai to handle loading in data for GANS. We will need a data loader that provides data for both the generator and the dicriminator.
bs = 128
size = 64
db = DataBlock(blocks = (TransformBlock, ImageBlock),
get_x = generate_noise,
get_items = get_image_files,
splitter = IndexSplitter([]),
item_tfms=Resize(size, method=ResizeMethod.Crop),
batch_tfms = Normalize.from_stats(torch.tensor([0.5,0.5,0.5]), torch.tensor([0.5,0.5,0.5])))
path=Path('/home/fast/.fastai/data/selfie2anime/trainA')
dls = db.dataloaders(path, path=path, bs=bs)
dls.show_batch(max_n=16)
As seen above, we need to provide examples of real images. These are given to the critic to teach it what real images look like. We will also feed it fake images from the generator during the training loop.
Below is an InvisibleTensor, that is fed to the generator. It is invisible and not shown when we show a batch, mostly because it has no real meaning to us, and is instead just a random number.
dls.one_batch()[0][0:3,0:20]
generator = basic_generator(64, n_channels=3, n_extra_layers=1)
critic = basic_critic (64, n_channels=3, n_extra_layers=1, act_cls=partial(nn.LeakyReLU, negative_slope=0.2))
learn = GANLearner.wgan(dls, generator, critic, opt_func = RMSProp)
We have to turn off validation metrics due to their not being a validation set for GANs, we also create a callback so the validation run stops early. We turn on training metrics so that we can track generator and critic loss seperately.
learn.recorder.train_metrics=True
learn.recorder.valid_metrics=False
with learn.no_logging():
with learn.no_bar():
learn.fit(50, 2e-4, wd=0.)
learn.recorder.plot_loss()
learn.show_results(max_n=9, ds_idx=0)
As you can see, the results are not the best right off the bat. You can get better results by training for longer, and finetuning the model more, but I think it is also important to show that GANs can take awhile to train.
with learn.no_logging():
with learn.no_bar():
learn.fit(500, 2e-4, wd=0.)
learn.recorder.plot_loss()
Here is what we get after 500 epochs. With this GAN I wouldn't expect great results. GANs in general have a lot of parameters and are hard to train. U-NET based GANS do fairly well for a low parameter Generators though. U-NETs are conditional Generators that take an image as an input. Feel free to look at other architectures as well, though know that they can take a long time to train.
learn.show_results(max_n=9, ds_idx=0)
And here is the generator model architecture. This is just like a CNN, except that it uses ConvTranspose2D
learn.model.generator
rand is going to be our input into our ConvTranspose layer. It is just random noise.
rand=torch.randn([1,100],device=default_device())
Here we are using an identity matrix * 5, we are going to add this to blow up specific weights. This is just to show example of how the random input gets mapped to different features.
ones = 5*torch.eye(100,device=default_device())
ones[0:5,0:5]
alter_single_params=(rand+ones)[0:10]
rand[:,0:10],alter_single_params[:,0:10]
for img in learn.model.generator(alter_single_params).clamp(0,1):
show_image(img)
Here we are slowing increasing the value of a single parameter, which slowly changes the output image in some way. There are other architectures that would have cleaner feature maps in the latent space than this one.
for i in range(10):
img=learn.model.generator(rand+ones[:,0]*0.3*(i-5))[0].clamp(0,1)
show_image(img)
def gen_loss(fake_pred, output, target): return fake_pred.mean()
def critic_loss(real_pred, fake_pred): return real_pred.mean() - fake_pred.mean()
This first module you will notice is AddChannels. This Module is simply responsible for simply making our noise into the correct shape for the following Conv Layer. Remember a ConvLayer expects a shape of BatchSize x Channels x Width x Height. We are going to reproduce this now. Later we will reproduce ConvTranspose using padding, and a regular Convolutional Layer.
learn.model.generator
class AddChannels(Module):
"Add `n_dim` channels at the end of the input."
def __init__(self, n_dim): self.n_dim=n_dim
def forward(self, x): return x.view(*(list(x.shape)+[1]*self.n_dim))
x=dls.one_batch()[0]
n_dim=2
list(x.shape),list(x.shape)+[1]*n_dim
x=AddChannels(2)(x)
x.shape
convT = nn.ConvTranspose2d(100, 512, kernel_size=(4, 4), stride=(1, 1), bias=False)
Below is the values we are trying to match using our padding+convultion approach.
convT(x.cpu())[0,0,0:4,0:4]
Here we are pad every parameter of our input with a lot of zeroes, this will end up giving us a single specific weight from our weights. We are effectively multiply everything by a bunch of zeroes, so this is not how this tends to be done in practice. Though it helps to understand how we are doing this, because instead of a ConvTranspose a lot of the time we instead use Reflection Padding and a regular Convolutional Layer.
x=F.pad(x, (3, 3, 3, 3), "constant", 0)
print(x[0][0])
print(x.shape)
F.conv2d(x.cpu(),convT.weight.transpose(0, 1).flip(2).flip(3))[0,0,0:4,0:4]
As you see these are the same, showing that we reemplemented ConvTRanspose in a similar, albeit inefficient way. Though, this gives us the intuitive understanding of why we did it this way.
If you would like to read more about GANs, then I would invite you to take a look at DeOldify(https://github.com/jantic/DeOldify). You can also take a look at last year's fastai GAN lesson https://course.fast.ai/videos/?lesson=7
export2html.notebook2html(fname='2020-06-11-Training-GANs.ipynb', dest='html/', template_file='fastpages.tpl',n_workers=1)