I need to build a model that takes the output of a transpose convolutional network and uses it as the mean parameter in a Normal distribution.
The code is as follows:
import torch
import torch.nn as nn
from torch.distributions import Normal
class Model(nn.Module):
def __init__(self):
super().__init__()
self.convT = nn.ConvTranspose2d(in_channels=3, out_channels=3, kernel_size=1)
self.conv = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=1)
def loss(self, x, permute=False, conv=False):
pi_mu = self.conv(x) if conv else self.convT(x)
pi_mu = pi_mu.permute(0, 2, 3, 1).unsqueeze(1) if permute else pi_mu
return torch.sum(Normal(pi_mu, torch.ones_like(pi_mu)).log_prob(torch.ones_like(pi_mu)))
model = Model()
loss = model.loss(torch.randn(8, 3, 64, 64))
loss.backward()
print('ConvTranspose and no permute/unsqueeze done.')
loss = model.loss(torch.randn(8, 3, 64, 64), permute=True, conv=True)
loss.backward()
print('Conv and permute/unsqueeze done.')
loss = model.loss(torch.randn(8, 3, 64, 64), permute=True)
loss.backward()
print('ConvTranspose and permute/unsqueeze done.')
I’ve outlined three variations in the code, the third of which raises an error:
- Use a ConvTranspose2D layer but do not permute the output. Works ok.
- Use a Conv2D later and permute the output. Works ok.
- Use a ConvTranspose2D and permute the output. Raises an error (see below).
The output is as follows (python3.7, torch1.4.0)
ConvTranspose and no permute/unsqueeze done.
Conv and permute/unsqueeze done.
Traceback (most recent call last):
File "/Users/alex/Developer/project/test.py", line 26, in <module>
loss.backward()
File "/Users/alex/miniconda3/envs/pytorch_env/lib/python3.7/site-packages/torch/tensor.py", line 195, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/Users/alex/miniconda3/envs/pytorch_env/lib/python3.7/site-packages/torch/autograd/__init__.py", line 99, in backward
allow_unreachable=True) # allow_unreachable flag
RuntimeError: columns needs to be contiguous