[MRG] Update example GAN to avoid the 10 minute CircleCI limit (#258)

* shortened example GAN

* pep8 and typo

* awesome animation

* small eror pep8

* add animation to doc

* better timing animation

* tune step

Author: rflamary

docs/source/conf.py

@@ -337,14 +337,16 @@ def __getattr__(cls, name):
 intersphinx_mapping = {'python': ('https://docs.python.org/3', None),
                        'numpy': ('http://docs.scipy.org/doc/numpy/', None),
                        'scipy': ('http://docs.scipy.org/doc/scipy/reference/', None),
-                       'matplotlib': ('http://matplotlib.org/', None)}
+                       'matplotlib': ('http://matplotlib.org/', None),
+                       'torch': ('https://pytorch.org/docs/stable/', None)}
 
 sphinx_gallery_conf = {
     'examples_dirs': ['../../examples', '../../examples/da'],
     'gallery_dirs': 'auto_examples',
     'backreferences_dir':  'gen_modules/backreferences',
 	'inspect_global_variables'  : True,
     'doc_module'          : ('ot','numpy','scipy','pylab'),
+    'matplotlib_animations': True,
     'reference_url': {
 	'ot': None}
 }

examples/backends/plot_wass2_gan_torch.py

@@ -50,6 +50,7 @@
 
 import numpy as np
 import matplotlib.pyplot as pl
+import matplotlib.animation as animation
 import torch
 from torch import nn
 import ot
@@ -112,10 +113,10 @@ def forward(self, x):
 
 
 G = Generator()
-optimizer = torch.optim.RMSprop(G.parameters(), lr=0.001)
+optimizer = torch.optim.RMSprop(G.parameters(), lr=0.00019, eps=1e-5)
 
 # number of iteration and size of the batches
-n_iter = 500
+n_iter = 200  # set to 200 for doc build but 1000 is better ;)
 size_batch = 500
 
 # generate statis samples to see their trajectory along training
@@ -167,7 +168,7 @@ def forward(self, x):
 
 pl.figure(3, (10, 10))
 
-ivisu = [0, 10, 50, 100, 150, 200, 300, 400, 499]
+ivisu = [0, 10, 25, 50, 75, 125, 15, 175, 199]
 
 for i in range(9):
     pl.subplot(3, 3, i + 1)
@@ -179,6 +180,37 @@ def forward(self, x):
     if i == 0:
         pl.legend()
 
+# %%
+# Animate trajectories of generated samples along iteration
+# -------------------------------------------------------
+
+pl.figure(4, (8, 8))
+
+
+def _update_plot(i):
+    pl.clf()
+    pl.scatter(xd[:, 0], xd[:, 1], label='Data samples from $\mu_d$', alpha=0.1)
+    pl.scatter(xvisu[i, :, 0], xvisu[i, :, 1], label='Data samples from $G\#\mu_n$', alpha=0.5)
+    pl.xticks(())
+    pl.yticks(())
+    pl.xlim((-1.5, 1.5))
+    pl.ylim((-1.5, 1.5))
+    pl.title('Iter. {}'.format(i))
+    return 1
+
+
+i = 0
+pl.scatter(xd[:, 0], xd[:, 1], label='Data samples from $\mu_d$', alpha=0.1)
+pl.scatter(xvisu[i, :, 0], xvisu[i, :, 1], label='Data samples from $G\#\mu_n$', alpha=0.5)
+pl.xticks(())
+pl.yticks(())
+pl.xlim((-1.5, 1.5))
+pl.ylim((-1.5, 1.5))
+pl.title('Iter. {}'.format(ivisu[i]))
+
+
+ani = animation.FuncAnimation(pl.gcf(), _update_plot, n_iter, interval=100, repeat_delay=2000)
+
 # %%
 # Generate and visualize data
 # ---------------------------
@@ -188,7 +220,7 @@ def forward(self, x):
 xn = torch.randn(size_batch, 2)
 x = G(xn).detach().numpy()
 
-pl.figure(4)
+pl.figure(5)
 pl.scatter(xd[:, 0], xd[:, 1], label='Data samples from $\mu_d$', alpha=0.5)
 pl.scatter(x[:, 0], x[:, 1], label='Data samples from $G\#\mu_n$', alpha=0.5)
 pl.title('Sources and Target distributions')