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()