We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
vae comparison example giving different results for Pyro and PyTorch implementations
Hi,

Pyro version: 3.3
PyTorch: 1.1

Running `examples/vae/vae_comparison.py`

With PyTorch:

```
downloading data
download complete.
downloading data
download complete.
Running PyTorch VAE implementation
====> Epoch: 0 
Training loss: 0.0022
Test set loss: 0.0027
====> Epoch: 1 
Training loss: 0.0021
Test set loss: 0.0024
====> Epoch: 2 
Training loss: 0.0021
Test set loss: 0.0022
====> Epoch: 3 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 4 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 5 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 6 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 7 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 8 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 9 
Training loss: 0.0021
Test set loss: 0.0021
```

With Pyro:

```
downloading data
download complete.
downloading data
download complete.
Running Pyro VAE implementation
====> Epoch: 0 
Training loss: 0.0017
Test set loss: 0.0013
====> Epoch: 1 
Training loss: 0.0012
Test set loss: 0.0012
====> Epoch: 2 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 3 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 4 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 5 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 6 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 7 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 8 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 9 
Training loss: 0.0011
Test set loss: 0.0011
```

Is this difference expected?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/vae/vae_comparison.py`
Content:
```
1 import argparse
2 import itertools
3 import os
4 from abc import ABCMeta, abstractmethod
5 
6 import torch
7 import torch.nn as nn
8 from six import add_metaclass
9 from torch.nn import functional
10 from torchvision.utils import save_image
11 
12 import pyro
13 from pyro.contrib.examples import util
14 from pyro.distributions import Bernoulli, Normal
15 from pyro.infer import SVI, JitTrace_ELBO, Trace_ELBO
16 from pyro.optim import Adam
17 from utils.mnist_cached import DATA_DIR, RESULTS_DIR
18 
19 """
20 Comparison of VAE implementation in PyTorch and Pyro. This example can be
21 used for profiling purposes.
22 
23 The PyTorch VAE example is taken (with minor modification) from pytorch/examples.
24 Source: https://github.com/pytorch/examples/tree/master/vae
25 """
26 
27 TRAIN = 'train'
28 TEST = 'test'
29 OUTPUT_DIR = RESULTS_DIR
30 
31 
32 # VAE encoder network
33 class Encoder(nn.Module):
34     def __init__(self):
35         super(Encoder, self).__init__()
36         self.fc1 = nn.Linear(784, 400)
37         self.fc21 = nn.Linear(400, 20)
38         self.fc22 = nn.Linear(400, 20)
39         self.relu = nn.ReLU()
40 
41     def forward(self, x):
42         x = x.reshape(-1, 784)
43         h1 = self.relu(self.fc1(x))
44         return self.fc21(h1), torch.exp(self.fc22(h1))
45 
46 
47 # VAE Decoder network
48 class Decoder(nn.Module):
49     def __init__(self):
50         super(Decoder, self).__init__()
51         self.fc3 = nn.Linear(20, 400)
52         self.fc4 = nn.Linear(400, 784)
53         self.relu = nn.ReLU()
54 
55     def forward(self, z):
56         h3 = self.relu(self.fc3(z))
57         return torch.sigmoid(self.fc4(h3))
58 
59 
60 @add_metaclass(ABCMeta)
61 class VAE(object):
62     """
63     Abstract class for the variational auto-encoder. The abstract method
64     for training the network is implemented by subclasses.
65     """
66 
67     def __init__(self, args, train_loader, test_loader):
68         self.args = args
69         self.vae_encoder = Encoder()
70         self.vae_decoder = Decoder()
71         self.train_loader = train_loader
72         self.test_loader = test_loader
73         self.mode = TRAIN
74 
75     def set_train(self, is_train=True):
76         if is_train:
77             self.mode = TRAIN
78             self.vae_encoder.train()
79             self.vae_decoder.train()
80         else:
81             self.mode = TEST
82             self.vae_encoder.eval()
83             self.vae_decoder.eval()
84 
85     @abstractmethod
86     def compute_loss_and_gradient(self, x):
87         """
88         Given a batch of data `x`, run the optimizer (backpropagate the gradient),
89         and return the computed loss.
90 
91         :param x: batch of data or a single datum (MNIST image).
92         :return: loss computed on the data batch.
93         """
94         return
95 
96     def model_eval(self, x):
97         """
98         Given a batch of data `x`, run it through the trained VAE network to get
99         the reconstructed image.
100 
101         :param x: batch of data or a single datum (MNIST image).
102         :return: reconstructed image, and the latent z's mean and variance.
103         """
104         z_mean, z_var = self.vae_encoder(x)
105         if self.mode == TRAIN:
106             z = Normal(z_mean, z_var.sqrt()).sample()
107         else:
108             z = z_mean
109         return self.vae_decoder(z), z_mean, z_var
110 
111     def train(self, epoch):
112         self.set_train(is_train=True)
113         train_loss = 0
114         for batch_idx, (x, _) in enumerate(self.train_loader):
115             loss = self.compute_loss_and_gradient(x)
116             train_loss += loss
117         print('====> Epoch: {} \nTraining loss: {:.4f}'.format(
118             epoch, train_loss / len(self.train_loader.dataset)))
119 
120     def test(self, epoch):
121         self.set_train(is_train=False)
122         test_loss = 0
123         for i, (x, _) in enumerate(self.test_loader):
124             with torch.no_grad():
125                 recon_x = self.model_eval(x)[0]
126                 test_loss += self.compute_loss_and_gradient(x)
127             if i == 0:
128                 n = min(x.size(0), 8)
129                 comparison = torch.cat([x[:n],
130                                         recon_x.reshape(self.args.batch_size, 1, 28, 28)[:n]])
131                 save_image(comparison.detach().cpu(),
132                            os.path.join(OUTPUT_DIR, 'reconstruction_' + str(epoch) + '.png'),
133                            nrow=n)
134 
135         test_loss /= len(self.test_loader.dataset)
136         print('Test set loss: {:.4f}'.format(test_loss))
137 
138 
139 class PyTorchVAEImpl(VAE):
140     """
141     Adapted from pytorch/examples.
142     Source: https://github.com/pytorch/examples/tree/master/vae
143     """
144 
145     def __init__(self, *args, **kwargs):
146         super(PyTorchVAEImpl, self).__init__(*args, **kwargs)
147         self.optimizer = self.initialize_optimizer(lr=1e-3)
148 
149     def compute_loss_and_gradient(self, x):
150         self.optimizer.zero_grad()
151         recon_x, z_mean, z_var = self.model_eval(x)
152         binary_cross_entropy = functional.binary_cross_entropy(recon_x, x.reshape(-1, 784))
153         # Uses analytical KL divergence expression for D_kl(q(z|x) || p(z))
154         # Refer to Appendix B from VAE paper:
155         # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
156         # (https://arxiv.org/abs/1312.6114)
157         kl_div = -0.5 * torch.sum(1 + z_var.log() - z_mean.pow(2) - z_var)
158         kl_div /= self.args.batch_size * 784
159         loss = binary_cross_entropy + kl_div
160         if self.mode == TRAIN:
161             loss.backward()
162             self.optimizer.step()
163         return loss.item()
164 
165     def initialize_optimizer(self, lr=1e-3):
166         model_params = itertools.chain(self.vae_encoder.parameters(), self.vae_decoder.parameters())
167         return torch.optim.Adam(model_params, lr)
168 
169 
170 class PyroVAEImpl(VAE):
171     """
172     Implementation of VAE using Pyro. Only the model and the guide specification
173     is needed to run the optimizer (the objective function does not need to be
174     specified as in the PyTorch implementation).
175     """
176 
177     def __init__(self, *args, **kwargs):
178         super(PyroVAEImpl, self).__init__(*args, **kwargs)
179         self.optimizer = self.initialize_optimizer(lr=1e-3)
180 
181     def model(self, data):
182         decoder = pyro.module('decoder', self.vae_decoder)
183         z_mean, z_std = torch.zeros([data.size(0), 20]), torch.ones([data.size(0), 20])
184         with pyro.plate('data', data.size(0)):
185             z = pyro.sample('latent', Normal(z_mean, z_std).to_event(1))
186             img = decoder.forward(z)
187             pyro.sample('obs',
188                         Bernoulli(img).to_event(1),
189                         obs=data.reshape(-1, 784))
190 
191     def guide(self, data):
192         encoder = pyro.module('encoder', self.vae_encoder)
193         with pyro.plate('data', data.size(0)):
194             z_mean, z_var = encoder.forward(data)
195             pyro.sample('latent', Normal(z_mean, z_var.sqrt()).to_event(1))
196 
197     def compute_loss_and_gradient(self, x):
198         if self.mode == TRAIN:
199             loss = self.optimizer.step(x)
200         else:
201             loss = self.optimizer.evaluate_loss(x)
202         loss /= self.args.batch_size * 784
203         return loss
204 
205     def initialize_optimizer(self, lr):
206         optimizer = Adam({'lr': lr})
207         elbo = JitTrace_ELBO() if self.args.jit else Trace_ELBO()
208         return SVI(self.model, self.guide, optimizer, loss=elbo)
209 
210 
211 def setup(args):
212     pyro.set_rng_seed(args.rng_seed)
213     train_loader = util.get_data_loader(dataset_name='MNIST',
214                                         data_dir=DATA_DIR,
215                                         batch_size=args.batch_size,
216                                         is_training_set=True,
217                                         shuffle=True)
218     test_loader = util.get_data_loader(dataset_name='MNIST',
219                                        data_dir=DATA_DIR,
220                                        batch_size=args.batch_size,
221                                        is_training_set=False,
222                                        shuffle=True)
223     global OUTPUT_DIR
224     OUTPUT_DIR = os.path.join(RESULTS_DIR, args.impl)
225     if not os.path.exists(OUTPUT_DIR):
226         os.makedirs(OUTPUT_DIR)
227     pyro.clear_param_store()
228     return train_loader, test_loader
229 
230 
231 def main(args):
232     train_loader, test_loader = setup(args)
233     if args.impl == 'pyro':
234         vae = PyroVAEImpl(args, train_loader, test_loader)
235         print('Running Pyro VAE implementation')
236     elif args.impl == 'pytorch':
237         vae = PyTorchVAEImpl(args, train_loader, test_loader)
238         print('Running PyTorch VAE implementation')
239     else:
240         raise ValueError('Incorrect implementation specified: {}'.format(args.impl))
241     for i in range(args.num_epochs):
242         vae.train(i)
243         if not args.skip_eval:
244             vae.test(i)
245 
246 
247 if __name__ == '__main__':
248     assert pyro.__version__.startswith('0.3.3')
249     parser = argparse.ArgumentParser(description='VAE using MNIST dataset')
250     parser.add_argument('-n', '--num-epochs', nargs='?', default=10, type=int)
251     parser.add_argument('--batch_size', nargs='?', default=128, type=int)
252     parser.add_argument('--rng_seed', nargs='?', default=0, type=int)
253     parser.add_argument('--impl', nargs='?', default='pyro', type=str)
254     parser.add_argument('--skip_eval', action='store_true')
255     parser.add_argument('--jit', action='store_true')
256     parser.set_defaults(skip_eval=False)
257     args = parser.parse_args()
258     main(args)
259
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/vae/vae_comparison.py b/examples/vae/vae_comparison.py
--- a/examples/vae/vae_comparison.py
+++ b/examples/vae/vae_comparison.py
@@ -103,7 +103,7 @@
         """
         z_mean, z_var = self.vae_encoder(x)
         if self.mode == TRAIN:
-            z = Normal(z_mean, z_var.sqrt()).sample()
+            z = Normal(z_mean, z_var.sqrt()).rsample()
         else:
             z = z_mean
         return self.vae_decoder(z), z_mean, z_var

{"golden_diff": "diff --git a/examples/vae/vae_comparison.py b/examples/vae/vae_comparison.py\n--- a/examples/vae/vae_comparison.py\n+++ b/examples/vae/vae_comparison.py\n@@ -103,7 +103,7 @@\n         \"\"\"\n         z_mean, z_var = self.vae_encoder(x)\n         if self.mode == TRAIN:\n-            z = Normal(z_mean, z_var.sqrt()).sample()\n+            z = Normal(z_mean, z_var.sqrt()).rsample()\n         else:\n             z = z_mean\n         return self.vae_decoder(z), z_mean, z_var\n", "issue": "vae comparison example giving different results for Pyro and PyTorch implementations\nHi,\r\n\r\nPyro version: 3.3\r\nPyTorch: 1.1\r\n\r\nRunning `examples/vae/vae_comparison.py`\r\n\r\nWith PyTorch:\r\n\r\n```\r\ndownloading data\r\ndownload complete.\r\ndownloading data\r\ndownload complete.\r\nRunning PyTorch VAE implementation\r\n====> Epoch: 0 \r\nTraining loss: 0.0022\r\nTest set loss: 0.0027\r\n====> Epoch: 1 \r\nTraining loss: 0.0021\r\nTest set loss: 0.0024\r\n====> Epoch: 2 \r\nTraining loss: 0.0021\r\nTest set loss: 0.0022\r\n====> Epoch: 3 \r\nTraining loss: 0.0021\r\nTest set loss: 0.0021\r\n====> Epoch: 4 \r\nTraining loss: 0.0021\r\nTest set loss: 0.0021\r\n====> Epoch: 5 \r\nTraining loss: 0.0021\r\nTest set loss: 0.0021\r\n====> Epoch: 6 \r\nTraining loss: 0.0021\r\nTest set loss: 0.0021\r\n====> Epoch: 7 \r\nTraining loss: 0.0021\r\nTest set loss: 0.0021\r\n====> Epoch: 8 \r\nTraining loss: 0.0021\r\nTest set loss: 0.0021\r\n====> Epoch: 9 \r\nTraining loss: 0.0021\r\nTest set loss: 0.0021\r\n```\r\n\r\nWith Pyro:\r\n\r\n```\r\ndownloading data\r\ndownload complete.\r\ndownloading data\r\ndownload complete.\r\nRunning Pyro VAE implementation\r\n====> Epoch: 0 \r\nTraining loss: 0.0017\r\nTest set loss: 0.0013\r\n====> Epoch: 1 \r\nTraining loss: 0.0012\r\nTest set loss: 0.0012\r\n====> Epoch: 2 \r\nTraining loss: 0.0011\r\nTest set loss: 0.0011\r\n====> Epoch: 3 \r\nTraining loss: 0.0011\r\nTest set loss: 0.0011\r\n====> Epoch: 4 \r\nTraining loss: 0.0011\r\nTest set loss: 0.0011\r\n====> Epoch: 5 \r\nTraining loss: 0.0011\r\nTest set loss: 0.0011\r\n====> Epoch: 6 \r\nTraining loss: 0.0011\r\nTest set loss: 0.0011\r\n====> Epoch: 7 \r\nTraining loss: 0.0011\r\nTest set loss: 0.0011\r\n====> Epoch: 8 \r\nTraining loss: 0.0011\r\nTest set loss: 0.0011\r\n====> Epoch: 9 \r\nTraining loss: 0.0011\r\nTest set loss: 0.0011\r\n```\r\n\r\nIs this difference expected?\n", "before_files": [{"content": "import argparse\nimport itertools\nimport os\nfrom abc import ABCMeta, abstractmethod\n\nimport torch\nimport torch.nn as nn\nfrom six import add_metaclass\nfrom torch.nn import functional\nfrom torchvision.utils import save_image\n\nimport pyro\nfrom pyro.contrib.examples import util\nfrom pyro.distributions import Bernoulli, Normal\nfrom pyro.infer import SVI, JitTrace_ELBO, Trace_ELBO\nfrom pyro.optim import Adam\nfrom utils.mnist_cached import DATA_DIR, RESULTS_DIR\n\n\"\"\"\nComparison of VAE implementation in PyTorch and Pyro. This example can be\nused for profiling purposes.\n\nThe PyTorch VAE example is taken (with minor modification) from pytorch/examples.\nSource: https://github.com/pytorch/examples/tree/master/vae\n\"\"\"\n\nTRAIN = 'train'\nTEST = 'test'\nOUTPUT_DIR = RESULTS_DIR\n\n\n# VAE encoder network\nclass Encoder(nn.Module):\n    def __init__(self):\n        super(Encoder, self).__init__()\n        self.fc1 = nn.Linear(784, 400)\n        self.fc21 = nn.Linear(400, 20)\n        self.fc22 = nn.Linear(400, 20)\n        self.relu = nn.ReLU()\n\n    def forward(self, x):\n        x = x.reshape(-1, 784)\n        h1 = self.relu(self.fc1(x))\n        return self.fc21(h1), torch.exp(self.fc22(h1))\n\n\n# VAE Decoder network\nclass Decoder(nn.Module):\n    def __init__(self):\n        super(Decoder, self).__init__()\n        self.fc3 = nn.Linear(20, 400)\n        self.fc4 = nn.Linear(400, 784)\n        self.relu = nn.ReLU()\n\n    def forward(self, z):\n        h3 = self.relu(self.fc3(z))\n        return torch.sigmoid(self.fc4(h3))\n\n\n@add_metaclass(ABCMeta)\nclass VAE(object):\n    \"\"\"\n    Abstract class for the variational auto-encoder. The abstract method\n    for training the network is implemented by subclasses.\n    \"\"\"\n\n    def __init__(self, args, train_loader, test_loader):\n        self.args = args\n        self.vae_encoder = Encoder()\n        self.vae_decoder = Decoder()\n        self.train_loader = train_loader\n        self.test_loader = test_loader\n        self.mode = TRAIN\n\n    def set_train(self, is_train=True):\n        if is_train:\n            self.mode = TRAIN\n            self.vae_encoder.train()\n            self.vae_decoder.train()\n        else:\n            self.mode = TEST\n            self.vae_encoder.eval()\n            self.vae_decoder.eval()\n\n    @abstractmethod\n    def compute_loss_and_gradient(self, x):\n        \"\"\"\n        Given a batch of data `x`, run the optimizer (backpropagate the gradient),\n        and return the computed loss.\n\n        :param x: batch of data or a single datum (MNIST image).\n        :return: loss computed on the data batch.\n        \"\"\"\n        return\n\n    def model_eval(self, x):\n        \"\"\"\n        Given a batch of data `x`, run it through the trained VAE network to get\n        the reconstructed image.\n\n        :param x: batch of data or a single datum (MNIST image).\n        :return: reconstructed image, and the latent z's mean and variance.\n        \"\"\"\n        z_mean, z_var = self.vae_encoder(x)\n        if self.mode == TRAIN:\n            z = Normal(z_mean, z_var.sqrt()).sample()\n        else:\n            z = z_mean\n        return self.vae_decoder(z), z_mean, z_var\n\n    def train(self, epoch):\n        self.set_train(is_train=True)\n        train_loss = 0\n        for batch_idx, (x, _) in enumerate(self.train_loader):\n            loss = self.compute_loss_and_gradient(x)\n            train_loss += loss\n        print('====> Epoch: {} \\nTraining loss: {:.4f}'.format(\n            epoch, train_loss / len(self.train_loader.dataset)))\n\n    def test(self, epoch):\n        self.set_train(is_train=False)\n        test_loss = 0\n        for i, (x, _) in enumerate(self.test_loader):\n            with torch.no_grad():\n                recon_x = self.model_eval(x)[0]\n                test_loss += self.compute_loss_and_gradient(x)\n            if i == 0:\n                n = min(x.size(0), 8)\n                comparison = torch.cat([x[:n],\n                                        recon_x.reshape(self.args.batch_size, 1, 28, 28)[:n]])\n                save_image(comparison.detach().cpu(),\n                           os.path.join(OUTPUT_DIR, 'reconstruction_' + str(epoch) + '.png'),\n                           nrow=n)\n\n        test_loss /= len(self.test_loader.dataset)\n        print('Test set loss: {:.4f}'.format(test_loss))\n\n\nclass PyTorchVAEImpl(VAE):\n    \"\"\"\n    Adapted from pytorch/examples.\n    Source: https://github.com/pytorch/examples/tree/master/vae\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        super(PyTorchVAEImpl, self).__init__(*args, **kwargs)\n        self.optimizer = self.initialize_optimizer(lr=1e-3)\n\n    def compute_loss_and_gradient(self, x):\n        self.optimizer.zero_grad()\n        recon_x, z_mean, z_var = self.model_eval(x)\n        binary_cross_entropy = functional.binary_cross_entropy(recon_x, x.reshape(-1, 784))\n        # Uses analytical KL divergence expression for D_kl(q(z|x) || p(z))\n        # Refer to Appendix B from VAE paper:\n        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014\n        # (https://arxiv.org/abs/1312.6114)\n        kl_div = -0.5 * torch.sum(1 + z_var.log() - z_mean.pow(2) - z_var)\n        kl_div /= self.args.batch_size * 784\n        loss = binary_cross_entropy + kl_div\n        if self.mode == TRAIN:\n            loss.backward()\n            self.optimizer.step()\n        return loss.item()\n\n    def initialize_optimizer(self, lr=1e-3):\n        model_params = itertools.chain(self.vae_encoder.parameters(), self.vae_decoder.parameters())\n        return torch.optim.Adam(model_params, lr)\n\n\nclass PyroVAEImpl(VAE):\n    \"\"\"\n    Implementation of VAE using Pyro. Only the model and the guide specification\n    is needed to run the optimizer (the objective function does not need to be\n    specified as in the PyTorch implementation).\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        super(PyroVAEImpl, self).__init__(*args, **kwargs)\n        self.optimizer = self.initialize_optimizer(lr=1e-3)\n\n    def model(self, data):\n        decoder = pyro.module('decoder', self.vae_decoder)\n        z_mean, z_std = torch.zeros([data.size(0), 20]), torch.ones([data.size(0), 20])\n        with pyro.plate('data', data.size(0)):\n            z = pyro.sample('latent', Normal(z_mean, z_std).to_event(1))\n            img = decoder.forward(z)\n            pyro.sample('obs',\n                        Bernoulli(img).to_event(1),\n                        obs=data.reshape(-1, 784))\n\n    def guide(self, data):\n        encoder = pyro.module('encoder', self.vae_encoder)\n        with pyro.plate('data', data.size(0)):\n            z_mean, z_var = encoder.forward(data)\n            pyro.sample('latent', Normal(z_mean, z_var.sqrt()).to_event(1))\n\n    def compute_loss_and_gradient(self, x):\n        if self.mode == TRAIN:\n            loss = self.optimizer.step(x)\n        else:\n            loss = self.optimizer.evaluate_loss(x)\n        loss /= self.args.batch_size * 784\n        return loss\n\n    def initialize_optimizer(self, lr):\n        optimizer = Adam({'lr': lr})\n        elbo = JitTrace_ELBO() if self.args.jit else Trace_ELBO()\n        return SVI(self.model, self.guide, optimizer, loss=elbo)\n\n\ndef setup(args):\n    pyro.set_rng_seed(args.rng_seed)\n    train_loader = util.get_data_loader(dataset_name='MNIST',\n                                        data_dir=DATA_DIR,\n                                        batch_size=args.batch_size,\n                                        is_training_set=True,\n                                        shuffle=True)\n    test_loader = util.get_data_loader(dataset_name='MNIST',\n                                       data_dir=DATA_DIR,\n                                       batch_size=args.batch_size,\n                                       is_training_set=False,\n                                       shuffle=True)\n    global OUTPUT_DIR\n    OUTPUT_DIR = os.path.join(RESULTS_DIR, args.impl)\n    if not os.path.exists(OUTPUT_DIR):\n        os.makedirs(OUTPUT_DIR)\n    pyro.clear_param_store()\n    return train_loader, test_loader\n\n\ndef main(args):\n    train_loader, test_loader = setup(args)\n    if args.impl == 'pyro':\n        vae = PyroVAEImpl(args, train_loader, test_loader)\n        print('Running Pyro VAE implementation')\n    elif args.impl == 'pytorch':\n        vae = PyTorchVAEImpl(args, train_loader, test_loader)\n        print('Running PyTorch VAE implementation')\n    else:\n        raise ValueError('Incorrect implementation specified: {}'.format(args.impl))\n    for i in range(args.num_epochs):\n        vae.train(i)\n        if not args.skip_eval:\n            vae.test(i)\n\n\nif __name__ == '__main__':\n    assert pyro.__version__.startswith('0.3.3')\n    parser = argparse.ArgumentParser(description='VAE using MNIST dataset')\n    parser.add_argument('-n', '--num-epochs', nargs='?', default=10, type=int)\n    parser.add_argument('--batch_size', nargs='?', default=128, type=int)\n    parser.add_argument('--rng_seed', nargs='?', default=0, type=int)\n    parser.add_argument('--impl', nargs='?', default='pyro', type=str)\n    parser.add_argument('--skip_eval', action='store_true')\n    parser.add_argument('--jit', action='store_true')\n    parser.set_defaults(skip_eval=False)\n    args = parser.parse_args()\n    main(args)\n", "path": "examples/vae/vae_comparison.py"}], "after_files": [{"content": "import argparse\nimport itertools\nimport os\nfrom abc import ABCMeta, abstractmethod\n\nimport torch\nimport torch.nn as nn\nfrom six import add_metaclass\nfrom torch.nn import functional\nfrom torchvision.utils import save_image\n\nimport pyro\nfrom pyro.contrib.examples import util\nfrom pyro.distributions import Bernoulli, Normal\nfrom pyro.infer import SVI, JitTrace_ELBO, Trace_ELBO\nfrom pyro.optim import Adam\nfrom utils.mnist_cached import DATA_DIR, RESULTS_DIR\n\n\"\"\"\nComparison of VAE implementation in PyTorch and Pyro. This example can be\nused for profiling purposes.\n\nThe PyTorch VAE example is taken (with minor modification) from pytorch/examples.\nSource: https://github.com/pytorch/examples/tree/master/vae\n\"\"\"\n\nTRAIN = 'train'\nTEST = 'test'\nOUTPUT_DIR = RESULTS_DIR\n\n\n# VAE encoder network\nclass Encoder(nn.Module):\n    def __init__(self):\n        super(Encoder, self).__init__()\n        self.fc1 = nn.Linear(784, 400)\n        self.fc21 = nn.Linear(400, 20)\n        self.fc22 = nn.Linear(400, 20)\n        self.relu = nn.ReLU()\n\n    def forward(self, x):\n        x = x.reshape(-1, 784)\n        h1 = self.relu(self.fc1(x))\n        return self.fc21(h1), torch.exp(self.fc22(h1))\n\n\n# VAE Decoder network\nclass Decoder(nn.Module):\n    def __init__(self):\n        super(Decoder, self).__init__()\n        self.fc3 = nn.Linear(20, 400)\n        self.fc4 = nn.Linear(400, 784)\n        self.relu = nn.ReLU()\n\n    def forward(self, z):\n        h3 = self.relu(self.fc3(z))\n        return torch.sigmoid(self.fc4(h3))\n\n\n@add_metaclass(ABCMeta)\nclass VAE(object):\n    \"\"\"\n    Abstract class for the variational auto-encoder. The abstract method\n    for training the network is implemented by subclasses.\n    \"\"\"\n\n    def __init__(self, args, train_loader, test_loader):\n        self.args = args\n        self.vae_encoder = Encoder()\n        self.vae_decoder = Decoder()\n        self.train_loader = train_loader\n        self.test_loader = test_loader\n        self.mode = TRAIN\n\n    def set_train(self, is_train=True):\n        if is_train:\n            self.mode = TRAIN\n            self.vae_encoder.train()\n            self.vae_decoder.train()\n        else:\n            self.mode = TEST\n            self.vae_encoder.eval()\n            self.vae_decoder.eval()\n\n    @abstractmethod\n    def compute_loss_and_gradient(self, x):\n        \"\"\"\n        Given a batch of data `x`, run the optimizer (backpropagate the gradient),\n        and return the computed loss.\n\n        :param x: batch of data or a single datum (MNIST image).\n        :return: loss computed on the data batch.\n        \"\"\"\n        return\n\n    def model_eval(self, x):\n        \"\"\"\n        Given a batch of data `x`, run it through the trained VAE network to get\n        the reconstructed image.\n\n        :param x: batch of data or a single datum (MNIST image).\n        :return: reconstructed image, and the latent z's mean and variance.\n        \"\"\"\n        z_mean, z_var = self.vae_encoder(x)\n        if self.mode == TRAIN:\n            z = Normal(z_mean, z_var.sqrt()).rsample()\n        else:\n            z = z_mean\n        return self.vae_decoder(z), z_mean, z_var\n\n    def train(self, epoch):\n        self.set_train(is_train=True)\n        train_loss = 0\n        for batch_idx, (x, _) in enumerate(self.train_loader):\n            loss = self.compute_loss_and_gradient(x)\n            train_loss += loss\n        print('====> Epoch: {} \\nTraining loss: {:.4f}'.format(\n            epoch, train_loss / len(self.train_loader.dataset)))\n\n    def test(self, epoch):\n        self.set_train(is_train=False)\n        test_loss = 0\n        for i, (x, _) in enumerate(self.test_loader):\n            with torch.no_grad():\n                recon_x = self.model_eval(x)[0]\n                test_loss += self.compute_loss_and_gradient(x)\n            if i == 0:\n                n = min(x.size(0), 8)\n                comparison = torch.cat([x[:n],\n                                        recon_x.reshape(self.args.batch_size, 1, 28, 28)[:n]])\n                save_image(comparison.detach().cpu(),\n                           os.path.join(OUTPUT_DIR, 'reconstruction_' + str(epoch) + '.png'),\n                           nrow=n)\n\n        test_loss /= len(self.test_loader.dataset)\n        print('Test set loss: {:.4f}'.format(test_loss))\n\n\nclass PyTorchVAEImpl(VAE):\n    \"\"\"\n    Adapted from pytorch/examples.\n    Source: https://github.com/pytorch/examples/tree/master/vae\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        super(PyTorchVAEImpl, self).__init__(*args, **kwargs)\n        self.optimizer = self.initialize_optimizer(lr=1e-3)\n\n    def compute_loss_and_gradient(self, x):\n        self.optimizer.zero_grad()\n        recon_x, z_mean, z_var = self.model_eval(x)\n        binary_cross_entropy = functional.binary_cross_entropy(recon_x, x.reshape(-1, 784))\n        # Uses analytical KL divergence expression for D_kl(q(z|x) || p(z))\n        # Refer to Appendix B from VAE paper:\n        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014\n        # (https://arxiv.org/abs/1312.6114)\n        kl_div = -0.5 * torch.sum(1 + z_var.log() - z_mean.pow(2) - z_var)\n        kl_div /= self.args.batch_size * 784\n        loss = binary_cross_entropy + kl_div\n        if self.mode == TRAIN:\n            loss.backward()\n            self.optimizer.step()\n        return loss.item()\n\n    def initialize_optimizer(self, lr=1e-3):\n        model_params = itertools.chain(self.vae_encoder.parameters(), self.vae_decoder.parameters())\n        return torch.optim.Adam(model_params, lr)\n\n\nclass PyroVAEImpl(VAE):\n    \"\"\"\n    Implementation of VAE using Pyro. Only the model and the guide specification\n    is needed to run the optimizer (the objective function does not need to be\n    specified as in the PyTorch implementation).\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        super(PyroVAEImpl, self).__init__(*args, **kwargs)\n        self.optimizer = self.initialize_optimizer(lr=1e-3)\n\n    def model(self, data):\n        decoder = pyro.module('decoder', self.vae_decoder)\n        z_mean, z_std = torch.zeros([data.size(0), 20]), torch.ones([data.size(0), 20])\n        with pyro.plate('data', data.size(0)):\n            z = pyro.sample('latent', Normal(z_mean, z_std).to_event(1))\n            img = decoder.forward(z)\n            pyro.sample('obs',\n                        Bernoulli(img).to_event(1),\n                        obs=data.reshape(-1, 784))\n\n    def guide(self, data):\n        encoder = pyro.module('encoder', self.vae_encoder)\n        with pyro.plate('data', data.size(0)):\n            z_mean, z_var = encoder.forward(data)\n            pyro.sample('latent', Normal(z_mean, z_var.sqrt()).to_event(1))\n\n    def compute_loss_and_gradient(self, x):\n        if self.mode == TRAIN:\n            loss = self.optimizer.step(x)\n        else:\n            loss = self.optimizer.evaluate_loss(x)\n        loss /= self.args.batch_size * 784\n        return loss\n\n    def initialize_optimizer(self, lr):\n        optimizer = Adam({'lr': lr})\n        elbo = JitTrace_ELBO() if self.args.jit else Trace_ELBO()\n        return SVI(self.model, self.guide, optimizer, loss=elbo)\n\n\ndef setup(args):\n    pyro.set_rng_seed(args.rng_seed)\n    train_loader = util.get_data_loader(dataset_name='MNIST',\n                                        data_dir=DATA_DIR,\n                                        batch_size=args.batch_size,\n                                        is_training_set=True,\n                                        shuffle=True)\n    test_loader = util.get_data_loader(dataset_name='MNIST',\n                                       data_dir=DATA_DIR,\n                                       batch_size=args.batch_size,\n                                       is_training_set=False,\n                                       shuffle=True)\n    global OUTPUT_DIR\n    OUTPUT_DIR = os.path.join(RESULTS_DIR, args.impl)\n    if not os.path.exists(OUTPUT_DIR):\n        os.makedirs(OUTPUT_DIR)\n    pyro.clear_param_store()\n    return train_loader, test_loader\n\n\ndef main(args):\n    train_loader, test_loader = setup(args)\n    if args.impl == 'pyro':\n        vae = PyroVAEImpl(args, train_loader, test_loader)\n        print('Running Pyro VAE implementation')\n    elif args.impl == 'pytorch':\n        vae = PyTorchVAEImpl(args, train_loader, test_loader)\n        print('Running PyTorch VAE implementation')\n    else:\n        raise ValueError('Incorrect implementation specified: {}'.format(args.impl))\n    for i in range(args.num_epochs):\n        vae.train(i)\n        if not args.skip_eval:\n            vae.test(i)\n\n\nif __name__ == '__main__':\n    assert pyro.__version__.startswith('0.3.3')\n    parser = argparse.ArgumentParser(description='VAE using MNIST dataset')\n    parser.add_argument('-n', '--num-epochs', nargs='?', default=10, type=int)\n    parser.add_argument('--batch_size', nargs='?', default=128, type=int)\n    parser.add_argument('--rng_seed', nargs='?', default=0, type=int)\n    parser.add_argument('--impl', nargs='?', default='pyro', type=str)\n    parser.add_argument('--skip_eval', action='store_true')\n    parser.add_argument('--jit', action='store_true')\n    parser.set_defaults(skip_eval=False)\n    args = parser.parse_args()\n    main(args)\n", "path": "examples/vae/vae_comparison.py"}]}