Fix dataparallel issue for draw (#13, #14)

Author: wohlert

draw/draw.py

@@ -154,6 +154,8 @@ def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1
         self.output = nn.Conv2d(in_channels, out_channels, **kwargs)
         self.state  = nn.Conv2d(in_channels, out_channels, **kwargs)
 
+        self.transform = nn.Conv2d(out_channels, in_channels, **kwargs)
+
     def forward(self, input, states):
         """
         Send input through the cell.
@@ -164,6 +166,8 @@ def forward(self, input, states):
         """
         (hidden, cell) = states
 
+        input = input + self.transform(hidden)
+
         forget_gate = torch.sigmoid(self.forget(input))
         input_gate  = torch.sigmoid(self.input(input))
         output_gate = torch.sigmoid(self.output(input))
@@ -280,4 +284,4 @@ def sample(self, x):
             h_dec, c_dec = self.decoder(torch.cat([z, canvas_next], dim=1), [h_dec, c_dec])
             canvas = canvas + F.pixel_shuffle(self.write_head(h_dec), 2)
 
-        return canvas
+        return canvas

run-draw.py

@@ -42,7 +42,7 @@
     loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, **kwargs)
 
     loss = nn.BCELoss(reduce=False).to(device)
-    
+
     for epoch in range(args.epochs):
         for x, _ in tqdm(loader):
             batch_size = x.size(0)
@@ -51,7 +51,7 @@
 
             x_hat, kl_divergence = model(x)
             x_hat = torch.sigmoid(x_hat)
-            
+
             reconstruction = loss(x_hat, x).sum(1)
             kl = kl_divergence.sum(1)
             elbo = torch.mean(reconstruction + kl)
@@ -67,10 +67,13 @@
                 print("Loss at step {}: {}".format(epoch, elbo.item()))
 
                 # Not sustainable if not dataparallel
-                x_sample = model.module.sample(args.batch_size)
+                if type(model) is nn.DataParallel:
+                    x_sample = model.module.sample(args.batch_size)
+                else:
+                    x_sample = model.sample(args.batch_size)
 
                 save_image(x_hat, "reconstruction-{}.jpg".format(epoch))
                 save_image(x_sample, "sample-{}.jpg".format(epoch))
 
             if epoch % 10 == 0:
-                torch.save(model, "model-{}.pt".format(epoch))
+                torch.save(model, "model-{}.pt".format(epoch))