Hi,
I am trying to recreate “PixelCNN” model written in Tensorflow at this link, in Pytroch.
My code is as below:
class MaskedConv2D(nn.Module):
def __init__(self, mask_type, kernel=5, filters=1):
super(MaskedConv2D, self).__init__()
self.kernel = kernel
self.filters = filters
self.mask_type = mask_type
self.weight = torch.nn.Parameter(torch.Tensor(self.kernel,self.kernel,1,self.filters), requires_grad=True)
self.bias = torch.nn.Parameter(torch.Tensor(self.filters,), requires_grad=True)
nn.init.xavier_uniform_(self.weight)
nn.init.constant_(self.bias, 0)
mask = np.ones(self.kernel**2, dtype=np.float32)
center = len(mask)//2
mask[center+1:] = 0
if self.mask_type == 'A':
mask[center] = 0
mask = mask.reshape((self.kernel, self.kernel, 1, 1))
self.mask = torch.tensor(mask, dtype=torch.float32, requires_grad=False)
# Keras code seem to pass "SAME to the padding parameter. How do I do it with Conv2d here?
def forward (self, inputs):
# If i try to do padding like below, I get "RuntimeError: Boolean value of Tensor with more than one value is ambiguous", during the Conv2d step.
inputs = F.pad(inputs, (1,1,0,0))
masked_w = self.weight.mul(self.mask)
output = torch.nn.Conv2d(in_channels=inputs,out_channels=64,kernel_size=masked_w, stride=1) # + self.bias
return torch.nn.ReLU(output)
class ResidualBlock(nn.Module):
def __init__(self, h=32):
super(ResidualBlock, self).__init__()
self.forward = nn.Sequential(MaskedConv2D('B', kernel=1, filters=h),
MaskedConv2D('B', kernel=3, filters=h),
MaskedConv2D('B', kernel=1, filters=2*h))
def call(self, inputs):
x = self.forward(inputs)
return x + inputs
class PixelCNN(nn.Module):
def __init__(self, hidden_features=64, output_features=64, resblocks_num=7):
super(PixelCNN, self).__init__()
self.model = nn.Sequential(
MaskedConv2D('A',kernel=7, filters=2*hidden_features),
ResidualBlock(hidden_features),
ResidualBlock(hidden_features),
ResidualBlock(hidden_features),
ResidualBlock(hidden_features),
ResidualBlock(hidden_features),
ResidualBlock(hidden_features),
ResidualBlock(hidden_features),
nn.Conv2d(in_channels = output_features, kernel_size = (1,1),out_channels=1),
torch.nn.ReLU(),
nn.Conv2d(in_channels = 1,out_channels = 1, kernel_size = (1,1)),
torch.nn.Sigmoid()
)
def forward(self, x):
return self.model(x)
net = PixelCNN()