diff --git a/src/layers/conv.jl b/src/layers/conv.jl index 338bb89725..f7c4ed2498 100644 --- a/src/layers/conv.jl +++ b/src/layers/conv.jl @@ -44,8 +44,25 @@ julia> layer3(xs) |> size # output size = `ceil(input_size/stride)` = 50 ``` """ struct SamePad end - -calc_padding(lt, pad, k::NTuple{N,T}, dilation, stride) where {T,N} = expand(Val(2*N), pad) +calc_padding(lt, pad::Int, k::NTuple{N,T}, dilation, stride) where {T,N} = expand(Val(2*N), pad) +function calc_padding(lt, pad::NTuple{Np, T}, k::NTuple{Nk, T}, dilation, stride) where {T, Nk, Np} + # calc_padding for when a tuple is passed as padding. + if Nk == Np + # duplicate each dim + new_pad = [] + for i in pad + push!(new_pad, i) + push!(new_pad, i) + end + return tuple(new_pad...) + elseif Nk == 2Np + # Copy as it is + return pad + else + # Error out + throw(ArgumentError("invalid padding dimensions")) + end +end function calc_padding(lt, ::SamePad, k::NTuple{N,T}, dilation, stride) where {N,T} #Ref: "A guide to convolution arithmetic for deep learning" https://arxiv.org/abs/1603.07285 diff --git a/test/layers/conv.jl b/test/layers/conv.jl index be26786495..f3dd6e8d3a 100644 --- a/test/layers/conv.jl +++ b/test/layers/conv.jl @@ -208,6 +208,16 @@ end m1 = ConvTranspose((3,5,3), 3=>6, stride=3) m2 = ConvTranspose((3,5,3), 3=>6, stride=3, outpad=(1,0,1)) @test size(m2(x))[1:3] == (size(m1(x))[1:3] .+ (1,0,1)) + + # test ConvTranspose constructor with tuple padding + kernel_dims = (3, 3) + pad = (1, 0) + x = randn(Float32, 5, 6, 1, 16) + m = ConvTranspose(kernel_dims, 1=>1, pad=pad) + result = m(x) + # Formula obtained from https://makeyourownneuralnetwork.blogspot.com/2020/02/calculating-output-size-of-convolutions.html + output_dims(pad_dims, kernel_dim, input_dim, stride) = (input_dim.-1).*stride.-(pad_dims.*2).+(kernel_dim.-1).+1 + @test output_dims(m.pad, kernel_dims, size(x)[1:2], m.stride) == size(result)[1:2] end @testset "CrossCor" begin