began.py

import torch as t
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import Adam

import numpy as np
import cv2
import os

batch_size = 16
img_size = 32
latent_dim = 64
channels = 1
lr = 0.0001 #0.0002
b1, b2 = 0.9, 0.999
n_epochs = 200
n_ch = 100  # Number of channels in initial convolution layers
sample_interval = 10  # Save a generated image every sample_interval number of batches

write_dir = '../images/began_gen_images/'
#data_dir = '../data/bub_single_24/'
data_dir = '../data/bub_single_24_shifted/'
os.makedirs(write_dir, exist_ok=True)

Zeros = t.zeros((batch_size, channels, img_size, img_size))
Ones = t.ones((batch_size, channels, img_size, img_size))

class Data(Dataset):
    def __init__(self):
        self.data = [cv2.imread(data_dir + im, cv2.IMREAD_GRAYSCALE)//255 for im in os.listdir(data_dir) if '-0.png' in im]
        self.data = t.Tensor([im for im in self.data if im.shape == (img_size, img_size)])
        #self.data += t.normal(mean=0.4, std=0.15, size=self.data.shape)
        #self.data = t.where(self.data > 1, t.ones_like(self.data), self.data)
        
        self.len = self.data.shape[0]

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        return self.data[index]

dataset = Data()

dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

class Generator(nn.Module):
    # Takes as input a random tensor of size (batch_size, latent_dim) generated by noise()
    def __init__(self):
        super(Generator, self).__init__()
        
        self.init_size = img_size//4
        self.linear = nn.Sequential(nn.Linear(latent_dim, n_ch * self.init_size**2))  #Outsize=(8x8xn) with n=128
        
        self.conv_blocks = nn.Sequential(
            nn.BatchNorm2d(n_ch),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(n_ch, n_ch, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(n_ch, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(n_ch, n_ch//2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(n_ch//2, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(n_ch//2, channels, kernel_size=3, stride=1, padding=1),
            #nn.Sigmoid()
            nn.Tanh()
            #Either using tanh or sigmoid and sending pixels < .5 to zero converges to bub on non-shifted data
            )
    
    def forward(self, noise):
        generated_im = self.linear(noise)
        generated_im = generated_im.view(generated_im.shape[0], n_ch, self.init_size, self.init_size)
        generated_im = self.conv_blocks(generated_im)
        #generated_im = t.where(generated_im < 0.45, Zeros, generated_im)
        return generated_im
        
    # (N, latent_dim) ->[linear] (N, n * 8**2) ->[reshape] (N, n, 8, 8) ->[upsampling] (N, n, 16, 16) ->[conv] (N, n, 16, 16) ->[upsampling] (N, n, 32, 32) ->[conv] (N, n/2, 32, 32) ->[conv] (N, 1, 32, 32)

class Discriminator(nn.Module):
    # Takes as input a data image or image generated by Generator
    def __init__(self):
        super(Discriminator, self).__init__()
        
        self.init_size = img_size//4
        self.block1 = nn.Sequential(
            nn.Conv2d(1, n_ch//2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(n_ch//2),
            nn.ReLU(),
            nn.Conv2d(n_ch//2, n_ch, kernel_size=3, stride=1, padding=1)
            )
        # Downsample
        self.block2 = nn.Sequential(
            nn.BatchNorm2d(n_ch),
            nn.ReLU(),
            nn.Conv2d(n_ch, n_ch, kernel_size=3, stride=1, padding=1)
            )
        # Downsample
        self.block3 = nn.Sequential(
            nn.BatchNorm2d(n_ch),
            nn.ReLU()
            )
            
        self.linear = nn.Sequential(
            nn.Linear(n_ch * self.init_size**2, img_size**2),
            nn.Sigmoid()
            #nn.Tanh()
            )
        
    def forward(self, input):
        input = self.block1(input)
        input = nn.functional.interpolate(input, size=(16, 16))
        input = self.block2(input)
        input = nn.functional.interpolate(input, size=(8, 8))
        input = self.block3(input)
        input = input.view(input.shape[0], n_ch * self.init_size**2)
        input = self.linear(input)
        input = input.view(input.shape[0], 1, img_size, img_size)
        #input = (input + 1)/2  #Scale the data from [-1,1] into [0,1]
        return input
        
    # (N, 1, 32, 32) ->[conv] (N, n//2, 32, 32) ->[conv] (N, n, 32, 32) ->[downsampling] (N, n, 16, 16) ->[conv] (N, n, 16, 16) ->[downsampling] (N, n, 8, 8) ->[reshape] (N, n * 8**2) ->[linear] (N, 32**2) ->[reshape] (N, 1, 32, 32)
    '''
    
        self.conv1 = nn.Sequential(nn.Conv2d(channels, n_ch//2, kernel_size=3, stride=2, padding=1), nn.ReLU())
        
        fc_dim = n_ch//2 * (img_size//2)**2
        self.fc = nn.Sequential(
            nn.Linear(fc_dim, 32),
            nn.BatchNorm1d(32, 0.8),
            nn.ReLU(inplace=True),
            nn.Linear(32, fc_dim),
            nn.BatchNorm1d(fc_dim),
            nn.ReLU(inplace=True)
            )
        
        self.up = nn.Sequential(nn.Upsample(scale_factor=2), nn.Conv2d(n_ch//2, channels, kernel_size=3, stride=1, padding=1))
    
    def forward(self, im):
        im = self.conv1(im)
        im = im.view(im.shape[0], -1)
        im = self.fc(im)
        im = im.view(im.shape[0], n_ch//2, img_size//2, img_size//2)
        return self.up(im)
    '''
    
def generate_noise():
    return t.rand((batch_size, latent_dim))

G = Generator()
D = Discriminator()

# Optimizers
optimizer_G = Adam(G.parameters(), lr=lr, betas=(b1, b2))
optimizer_D = Adam(D.parameters(), lr=lr, betas=(b1, b2))

# BEGAN hyper parameters
gamma = 0.75 # gamma in [0,1] controls image diversity.  Higher values of gamma lead to more diverse images
lambda_k = 0.001  # Learning rate for k
k = 0.0

def Loss(x):
    return t.mean(t.abs(x - D(x)))

for epoch in range(n_epochs):
    for i, data_imgs in enumerate(dataloader):
        
        # Configure generator input
        data_shape = data_imgs.shape
        data_imgs = data_imgs.view(data_imgs.shape[0], channels, data_imgs.shape[1], data_imgs.shape[2])
        
        # ----------------
        # Train Generator
        # ----------------
        
        optimizer_G.zero_grad()
        
        # Sample noise as generator input
        noise = generate_noise()
        
        # Generate a batch of images
        gen_imgs = G(noise)
        
        # Loss measures generator's ability to fool the discriminator
        G_loss = Loss(gen_imgs)

        G_loss.backward()
        optimizer_G.step()
        
        # --------------------
        # Train Discriminator
        # --------------------
        
        optimizer_D.zero_grad()
        
        # Measure discriminator's ability to classify real from generated samples
        D_loss_real = Loss(data_imgs)
        D_loss_fake = Loss(gen_imgs.detach())
        
        D_loss = D_loss_real - k*D_loss_fake
        
        D_loss.backward()
        optimizer_D.step()
        
        # ----------------
        # Update Weights
        # ----------------

        diff = t.mean(gamma*D_loss_real - D_loss_fake)

        # Update weight term for fake samples
        k = k + lambda_k*diff.item()
        k = min(max(k, 0), 1)  # Constraint to interval [0, 1]

        # Update convergence metric
        M = (D_loss_real + t.abs(diff)).item()
        
        # ----------------
        # Log Progress
        # ----------------
        
        print(
            "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] -- M: %f, k: %f"
            % (epoch, n_epochs, i, len(dataloader), D_loss.item(), G_loss.item(), M, k)
        )

        batches_done = epoch*len(dataloader) + i
        if batches_done % sample_interval == 0:
            os.makedirs(write_dir + str(batches_done), exist_ok=True)
            for j, gen_im in enumerate(gen_imgs):
                cv2.imwrite('%s%d/%d.png' % (write_dir, batches_done, j), gen_im[0].detach().numpy()*255)