Hi all,
I am tying to implement a simple 2d causal convolution based PixelCNN. My PyTorch implementation easily reach a sensible per-sample loss around 90+/- but my Julia’s Flux based solution stuck at 200+ loss per sample.
I initially thought the problem must be caused by the raster-scan order and column major mismatch. But it is counter-intuitive to think flipping the image by W and H would cause the autoregressive performance deteriorates drastically suddenly. Scanning an image column wisely should be a good order.
The 2D causal convolution can be implemented by multiply the weight of the kernel with a mask matrix. The mask is of the same size as the weight, but it masks out all “future” weight. The following code produce the two types of masks. The only difference lies in whether the centre pixel can be seen or not.
function create_mask(conv::Conv, type::Symbol=:B)
mask = similar(conv.weight)
mask .= 0 ## initialise all to zero
kh, kw = size(mask)[1:2]
@assert isodd(kw) && isodd(kh)
central_idx1 = kh ÷ 2 + 1
central_idx2 = kw ÷ 2 + 1
mask[1:(central_idx1-1), :, :, :] .= 1
mask[central_idx1, 1:central_idx2, :, :] .= 1
if type == :A
mask[central_idx1, central_idx2, :, :] .= 0
end
return mask
end
Then CausalConv is pretty standard. The forward pass masks the weight, then use the low level conv to do the difficult bit. One thing to note is: flipped = true, i.e. we match the default PyTorch behaviour where the weight is not flipped before the convolution.
struct MaskedConv{C<:Conv}
conv::C
mask
end
Flux.@layer MaskedConv trainable = (conv,)
function MaskedConv(k_size::Tuple{Int,Int}, ch::Pair{Int,Int}, act::Function = identity; mask_type::Symbol=:B, kwargs...)
ch_in, ch_out = ch
k_w, k_h = k_size
conv = Conv(k_size, ch, act; pad=:same, kwargs...)
mask = create_mask(conv, mask_type)
MaskedConv(conv, mask)
end
function (m::MaskedConv)(x)
# This creates a TEMPORARY masked tensor for this forward pass.
masked_weight = m.conv.weight .* m.mask
# Pass the temporary matrix into the convolution
(; pad, stride, dilation, groups) = m.conv
out = Flux.conv(x, masked_weight; flipped = true, pad, stride, dilation, groups)
# 4. Apply the bias block and the activation function smoothly
if !(typeof(m.conv.bias) == Bool && m.conv.bias == false)
out = out .+ reshape(m.conv.bias, 1, 1, :, 1)
end
return m.conv.σ.(out)
end
Then we can link a sequence of the causal layers together to create a neural net:
struct PixelCNN2D{T<:Chain}
d::Tuple{Int,Int} ## number of input; or input dimension
in_ch::Int
net::T ## trainable network
end
Flux.@layer PixelCNN2D trainable = (net,)
function PixelCNN2D(d::Tuple{Int,Int}=(28, 28), ch_size::Int=64, k_size::Tuple{Int,Int}=(7, 7), num_layers::Int=4, input_ch::Int=1, output_ch::Int=1; act=relu, bias=false, kwargs...)
net = Any[MaskedConv(k_size, input_ch=>ch_size, act; mask_type=:A, bias=bias, kwargs...)]
for _ in 1:(num_layers-1)
push!(net, MaskedConv(k_size, ch_size=>ch_size, act; mask_type=:B, bias=bias, kwargs...))
end
push!(net, MaskedConv(k_size, ch_size=>output_ch, identity; mask_type=:B, bias=bias, kwargs...))
net = Chain(net...)
return PixelCNN2D(d, input_ch, net)
end
function (m::PixelCNN2D)(xs)
m.net(xs)
end
I think my implementation matches the PyTorch almost verbatim. But my Julia implementation fails to learn.
The training steps are standard as well
model = PixelCNN2D() |> device
opt = Flux.setup(Adam(), model)
epochs = 20
for epoch in 1:epochs
for (images,) in mnist_binary_loader
batch_size = size(images)[end]
images = permutedims(images, (2, 1, 3, 4)) ## try to mimic the same behaviour as PyTorch, removing it doesnt help either
images = images |> device
li, gs = Flux.withgradient(model) do model
logits = model(images)
loss = Flux.logitbinarycrossentropy(
logits, images; agg=sum) / batch_size
end
Flux.update!(opt, model, gs[1])
end
end
I have tried the following:
- set Flux.conv’s flipped = false , which is the default behaviour of Julia’s convolution, then flip the mask accordingly to match, which doesn’t work
- remove all mask, the loss easily goes to near 0.0 which shows the masking making the problem too challenging
- permutedim images/or not, the result remains the same.
I am really stuck here and of course I can just use PyTorch. But I have already have most of the autoregressive models implemented in Flux and Julia and they all working fine. And I really wanted to know what is wrong.
To make it complete, the code below can be used to download and load the binary data in case anyone wants to try the code above. The PyTorch implementation is also listed below.
function load_binary_mnist_data(batch_size::Int=2000; tensor_shape=(28^2, :), seed=2345, train_size=50000, url="https://github.com/mgermain/MADE/releases/download/ICML2015/binarized_mnist.npz")
Random.seed!(seed)
# bmnist = npzread(download(url))
bmnist = npzread("./datasets/binarized_mnist.npz")
reshape_tf = x -> reshape(x, tensor_shape)
mnist_binary_test = BitMatrix(bmnist["test_data"]') |> reshape_tf
mnist_binary_train = BitMatrix(bmnist["train_data"]') |> reshape_tf
mnist_binary_valid = BitMatrix(bmnist["valid_data"]') |> reshape_tf
train_size = max(train_size, size(mnist_binary_train)[end])
train_data = selectdim(mnist_binary_train, ndims(mnist_binary_train), 1:train_size)
train_loader = Flux.DataLoader((Float32.(train_data),), batchsize=batch_size, shuffle=true)
test_loader = Flux.DataLoader((Float32.(mnist_binary_test),), batchsize=batch_size, shuffle=true)
valid_loader = Flux.DataLoader((Float32.(mnist_binary_valid),), batchsize=batch_size, shuffle=true)
return train_loader, test_loader, valid_loader
end
class MaskConv2d(nn.Conv2d):
def __init__(self, mask_type, *args, **kwargs):
assert mask_type == 'A' or mask_type == 'B'
super().__init__(*args, **kwargs)
self.register_buffer('mask', torch.zeros_like(self.weight))
self.create_mask(mask_type)
def forward(self, input):
return F.conv2d(input, self.weight * self.mask, self.bias, self.stride,
self.padding, self.dilation, self.groups)
def create_mask(self, mask_type):
k = self.kernel_size[0]
self.mask[:, :, :k // 2] = 1
self.mask[:, :, k // 2, :k // 2] = 1
if mask_type == 'B':
self.mask[:, :, k // 2, k // 2] = 1
class PixelCNN(nn.Module):
def __init__(self, device):
super().__init__()
model = [MaskConv2d('A', 1, 64, 7, padding=3), nn.ReLU()]
for _ in range(3):
model.extend([MaskConv2d('B', 64, 64, 7, padding=3), nn.ReLU()])
model.append(MaskConv2d('B', 64, 1, 7,padding=3))
self.net = nn.Sequential(*model)
self.device = device
def forward(self, x):
return self.net(x)
Any help is greatly appreciated!! Thanks again.