*Work In Progress: Some major code refactoring is underway and this branch is unstable a.t.m *.

Unsupervised representation learning with CapsNet based Generative Adversarial Networks

Arxiv paper link here

Coalesces Wassertein Distance algorithm based generative adversarial networks with deep convolutinal network based discriminator replaced by CapsNet architecture.

• Python > 3.6
• PyTorch
• TorchVision
• TorchNet
• TQDM
• Visdom [optional]

** For training, an NVIDIA TITAN XP GPU was used. Using CUDA with GPU is strongly recommended. CPU is supported but speed of training will be very slow.

• MNIST
• CIFAR-10
• SmallNORB

Running MNIST

Pass in the gpu device number for e.g. 0

$ python --dataset mnist --nc 1 --imageSize 32 --dataroot ./data --cuda {GPU_DEVICE_NUMBER}  --workers 4

Architecture

High level architecture:

Running SmallNORB

SmallNORB is an ideal dataset for testing the efficacy of Capsule networks over traditional CNN based GANs.

Pass in the gpu device number for e.g. 0

$ python --dataset smallnorb --nc 1 --imageSize 64 --dataroot ./data/smallnorb/ --cuda {GPU_DEVICE_NUMBER}  --workers 4

SmallNORB Results

Reconstructed smallNORB cars and planes:

Running CIFAR10

Pass in the gpu device number for e.g. 0

$ python3 main.py --nc 3 --dataset cifar10 --dataroot ./data --cuda {GPU_DEVICE_NUMBER} --workers 4 --niter [NUM_EPOCHS] 

Utilizes FAIR's visdom as visulization tool. If you'd like to visualize the test and train results, run with visualize args.

$ sudo python3 -m visdom.server &
$ python3 main.py --visualize --cuda

To run with MLP as G or D, run:

$ python3 main.py --dataset cifar10 --dataroot ./data --cuda {device_num} --experiment {Name} --mlp_G --ngf 512

Using CUDA

Pass in the gpu device number for e.g. 0

$ python main.py --cuda {GPU_DEVICE_NUMBER}

Enable Visualization

Start the server (probably in a screen or tmux):

python -m visdom.server -port 8097

Run with --visualize parameter

$ python main.py --cude {GPU_DEVICE_NUMBER} --visualize

Using DCGAN (and variants - BN, no-BN) as baseline against the CapsNet architecture.

Please send an email to raeidsaqur[at]cs[dot]toronto[dot]edu for questions, PRs etc.

*** Note: Improved ReadMe is in the works!