hinofafa / self-attention-hearthstone-gan

This repository provides a PyTorch implementation of SAGAN cited by heykeetae/Self-Attention-GAN. This repository provide an efficient method to generate large resolution images and attention weights visualisation using tensorboard platform. Tensorboard is a robust platform to monitor generated images and learning weights in computer vision learning experiment.

Self-Attention GAN in Hearthstone

**[Han Zhang, Ian Goodfellow, Dimitris Metaxas and Augustus Odena, "Self-Attention Generative Adversarial Networks." arXiv preprint arXiv:1805.08318 (2018)](https://arxiv.org/abs/1805.08318).**

Meta overview

This repository provides a PyTorch implementation of SAGAN. Both wgan-gp and wgan-hinge loss are ready, but note that wgan-gp is somehow not compatible with the spectral normalization. Remove all the spectral normalization at the model for the adoption of wgan-gp. Self-attentions are applied before CNN of both discriminator and generator.

Self Attention Layer

Original Repo status

• Unsupervised setting (use no label yet)
• Applied: Spectral Normalization, code from here
• Implemented: self-attention module, two-timescale update rule (TTUR), wgan-hinge loss, wgan-gp loss

Current Repo status

• Parallel Computation on multi-GPU
• Tensorboard loggings
• Attention visualization on 64 * 64 image
• Create Attention map of 64 * 64 image (4096 * 4096)
• Change custom (hearthstone) dataset
• Create 256*256 image [branch pix256]

Warning: 64*64 is the maximum 2power size of attention map for training in 2 Nvidia GTX 1080 Ti (24GB RAM)

Prerequisites

• Python 3.5+
• PyTorch 0.3.0
• opencv-python
• Details in requirements.txt

Usage

1. Clone the repository

$ git clone https://github.com/heykeetae/Self-Attention-GAN.git
$ cd Self-Attention-GAN
# for conda user
$ conda create -n sagan python=3.5
$ conda activate sagan
$ conda install pytorch=0.3.0

$ pip install -r requirements.txt

2. Install datasets (CelebA or LSUN or Hearthstone)

$ cd data
$ bash download.sh CelebA (404 not found)
# or
$ bash download.sh LSUN
# For Hearthstone player
$ mkdir hearthstone-card-images
$ cd hearthstone-card-images
$ wget https://www.dropbox.com/s/vvaxb4maoj4ri34/hearthstone_card.zip?dl=0
$ unzip hearthstone_card.zip?dl=0

3. Train

(i) Train in CelebA or Sagan dataset

$ python main.py --batch_size 64 --imsize 64 --dataset celeb --adv_loss hinge --version sagan_celeb
# or
$ python main.py --batch_size 64 --imsize 64 --dataset lsun --adv_loss hinge --version sagan_lsun

(ii) Custom parameteric Train in Hearthstone dataset

$ python main.py --batch_size 16 --imsize 64 --dataset hearthstone --adv_loss hinge --version sagan_hearth_at1 --num_workers 16 --use_tensorboard True --parallel True --total_step 100000 --log_step 100

For argument details, please read parameter.py

4. Attention & Statistics visualization

tensorboard --logdir ./logs/sagan_hearth_at1

5. Fake images located at

$ cd samples/sagan_celeb
# or
$ cd samples/sagan_lsun
# or
$ cd samples/sagan_hearth_at1

Samples generated every 100 iterations are located. The rate of sampling could be controlled via --sample_step (ex, --sample_step 100).

64*64 Results (step #95500)

64*64 Attention result on Hearthstone (step #95500)

• Colormap from opencv(https://docs.opencv.org/2.4/modules/contrib/doc/facerec/colormaps.html)
• Most attent part shows in RED (1) , most non-attent part shows in BLUE(0)
• Scores are ranged in [0,1]:

 

self-attention-hearthstone-gan's Issues

Large pixel fake image size

Create a new branch for create a large 256 * 256 hearthstone image.

Tensorboard Setting has inserted into the project

Tensorboard feature has inserted to the repository.

add a tensorboard logging to the train.py by inserting the following code at line 180. This needs to be modified because l4 does not exist if imsize is less than 64 and the network fails if imsize is greater than 64. When the larger imsize is fixed I will issue a pull request. Also build_tensorboard should be changed. Hope this helps!

def build_tensorboard(self):
    from logger import Logger
    if os.path.exists(self.log_path):
        shutil.rmtree(self.log_path)
    os.makedirs(self.log_path)
    self.logger = Logger(self.log_path)

Insert at line 180 in train.py

Print out log info

if (step + 1) % self.log_step == 0:
elapsed = time.time() - start_time
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Elapsed [{}], G_step [{}/{}], D_step[{}/{}], d_out_real: {:.4f}, "
" ave_gamma_l3: {:.4f}, ave_gamma_l4: {:.4f}".
format(elapsed, step + 1, self.total_step, (step + 1),
self.total_step, d_loss_real.item(),
self.G.attn1.gamma.mean().item(), self.G.attn2.gamma.mean().item()))

            # (1) Log values of the losses (scalars)
            info = {
                'd_loss_real': d_loss_real.item(),
                'd_loss_fake': d_loss_fake.item(),
                'd_loss': d_loss.item(),
                'g_loss_fake': g_loss_fake.item(),
                'ave_gamma_l3': self.G.attn1.gamma.mean().item(),
                'ave_gamma_l4': self.G.attn2.gamma.mean().item(),
            }


            for tag, value in info.items():
                self.logger.scalar_summary(tag, value, step + 1)


        # Sample images / Save and log
        if (step + 1) % self.sample_step == 0:

            # (2) Log values and gradients of the parameters (histogram)
            for net, name in zip([self.G, self.D], ['G_', 'D_']):
                for tag, value in net.named_parameters():
                    tag = name + tag.replace('.', '/')
                    self.logger.histo_summary(tag, self.to_np(value), step + 1)

            # (3) Log the images
            info = {

                'fake_images': self.to_np(fake_images.view(*display_vars)[:10, :, :, :]),
                'real_images': self.to_np(real_images.view(*display_vars)[:10, :, :, :]),
            }

            fake_images, _, _ = self.G(fixed_z)
            save_image(denorm(fake_images.data),
                       os.path.join(self.sample_path, '{}_fake.png'.format(step + 1)))

            info['fixed_fake_images'] = self.to_np(denorm(real_images.data).view(*display_vars)[:10, :, :, :])

            for tag, image in info.items():
                self.logger.image_summary(tag, image, step + 1)