Add fold and unfold

src/NNlib.jl

@@ -61,6 +61,9 @@ export conv, conv!, ∇conv_data, ∇conv_data!, ∇conv_filter,
 include("conv_bias_act.jl")
 export conv_bias_act, conv_bias_act!
 
+include("fold.jl")
+export unfold, unfold!, fold, fold!
+
 include("ctc.jl")
 export ctc_loss
 

src/fold.jl

@@ -0,0 +1,137 @@
+
+"""
+ unfold(x, kernel_size; stride = 1, pad = 0, dilation = 0, flipped = false)
+
+Places sliding windows of x into a container tensor of size (num_windows, window_size, batchsize).
+The window size is determined by the prod(spatial dims of kernel)*input_channels.
+The number of sliding windows will match those of convolution (conv) with the same kernel_size and arguments.
+Uses NNlib.im2col! as backend.
+"""
+function unfold(x::AbstractArray{T, N}, kernel_size::NTuple{K}; stride = 1, pad = 0, dilation = 1, flipped = false) where {T, K, N}
+ stride = expand(Val(N - 2), stride)
+ padding = expand(Val(N - 2), pad)
+ dilation = expand(Val(N - 2), dilation)
+ cdims = DenseConvDims(size(x), kernel_size; stride, padding, dilation, flipkernel=flipped)
+ return unfold(x, cdims)
+end
+
+"""
+ fold(y, output_size, kernel_size; stride = 1, pad = 0, dilation = 0, flipped = false)
+
+Accumulates sliding windows from the output of unfold into a container tensor of size `output_size`.
+An inverse to `unfold` can be obtained by using `fold` and accounting for scaling issues. 
+For example,
+
+```jldoctest
+julia> kernel_size, pad = (3, 3, 1, 1), 1;
+
+julia> x = reshape(1:64, 8, 8, 1, 1) |> collect;
+
+julia> y = unfold(x, kernel_size; pad=pad);
+
+julia> size(y)
+(64, 9, 1)
+
+julia> z = fold(y, size(x), kernel_size; pad=pad);
+
+julia> d = fold(unfold(ones(eltype(x), size(x)...), kernel_size; pad=pad), size(x), kernel_size; pad=pad)
+8×8×1×1 Array{Int64, 4}:
+[:, :, 1, 1] =
+ 4 6 6 6 6 6 6 4
+ 6 9 9 9 9 9 9 6
+ 6 9 9 9 9 9 9 6
+ 6 9 9 9 9 9 9 6
+ 6 9 9 9 9 9 9 6
+ 6 9 9 9 9 9 9 6
+ 6 9 9 9 9 9 9 6
+ 4 6 6 6 6 6 6 4
+
+julia> x == z./d
+true
+
+```
+Uses NNlib.col2im! as backend.
+"""
+function fold(x::AbstractArray{T, 3}, output_size::NTuple{N}, kernel_size::NTuple{K}; stride = 1, pad = 0, dilation = 1, flipped = false) where {T, K, N}
+ stride = expand(Val(N - 2), stride)
+ padding = expand(Val(N - 2), pad)
+ dilation = expand(Val(N - 2), dilation)
+ cdims = DenseConvDims(output_size, kernel_size; stride, padding, dilation, flipkernel=flipped)
+ return fold(x, output_size, cdims)
+end
+
+# im2col_dims returns (numblocks, blocksize, threadnum) where thread dim is used as thread-local
+# workspace for multithreaded conv. Ultimately, we want to threadnum with batchsize.
+unfold_dims(cdims::DenseConvDims) = im2col_dims(cdims)[1:2]
+
+# auto-allocating versions
+function unfold(x::AbstractArray{T, N}, cdims::DenseConvDims) where {T, N}
+ y = similar(x, unfold_dims(cdims)..., size(x, N)) # (numblocks, blocksize, batchsize)
+ return unfold!(y, x, cdims)
+end
+
+function fold(y::AbstractArray{T, 3}, output_size::NTuple, cdims::DenseConvDims) where {T}
+ x = similar(y, output_size) 
+ return fold!(x, y, cdims)
+end
+
+# N < 5 -dimension in-place versions 
+function unfold!(y::AbstractArray{yT, 3}, x::AbstractArray{xT, N}, cdims::DenseConvDims) where {yT, xT, N}
+ unfold!(
+ y, 
+ insert_singleton_spatial_dimension(x, 5-N), 
+ insert_singleton_spatial_dimension(cdims, 5-N), 
+ )
+ return y
+end
+
+function fold!(x::AbstractArray{xT, N}, y::AbstractArray{yT, 3}, cdims::DenseConvDims) where {yT, xT, N}
+ fold!(
+ insert_singleton_spatial_dimension(x, 5-N), 
+ y,
+ insert_singleton_spatial_dimension(cdims, 5-N), 
+ )
+ return x
+end
+
+# 5-dimension in-place versions 
+function unfold!(y::AbstractArray{yT, 3}, x::AbstractArray{xT, 5}, cdims::DenseConvDims) where {yT, xT}
+ @threads for batch_idx in 1:size(x, 5)
+ y_slice = view(y, :, :, batch_idx)
+ im2col!(y_slice, view(x, :, :, :, :, batch_idx), cdims)
+ end
+ return y
+end
+
+function fold!(x::AbstractArray{xT, 5}, y::AbstractArray{yT, 3}, cdims::DenseConvDims) where {xT, yT}
+ @threads for batch_idx in 1:size(x, 5)
+ y_slice = view(y, :, :, batch_idx)
+ col2im!(view(x, :, :, :, :, batch_idx), y_slice, cdims)
+ end
+ return x
+end
+
+# reverse diff rules
+function rrule(::typeof(unfold), x, cdims::DenseConvDims; kw...)
+ function unfold_pullback(Δ)
+ return (
+ NoTangent(),
+ fold(unthunk(Δ), size(x), cdims; kw...),
+ NoTangent(),
+ )
+ end
+ return unfold(x, cdims; kw...), unfold_pullback
+end
+
+function rrule(::typeof(fold), x, output_size, cdims::DenseConvDims; kw...)
+ function fold_pullback(Δ)
+ return (
+ NoTangent(),
+ unfold(unthunk(Δ), cdims; kw...),
+ NoTangent(),
+ NoTangent(),
+ )
+ end
+ return fold(x, output_size, cdims; kw...), fold_pullback
+end
+

test/fold.jl

@@ -0,0 +1,38 @@
+using NNlib, Test
+
+@testset "unfold wrapper" begin
+ x = rand(rng, 16, 16, 3, 10)
+ w = rand(rng, 5, 5, 3, 2)
+ @test size(unfold(x, size(w))) == (144, 75, 10)
+ @test size(unfold(x, size(w); pad=2)) == (256, 75, 10)
+ @test size(unfold(x, size(w); stride=2)) == (36, 75, 10)
+ @test size(unfold(x, size(w); dilation=2)) == (64, 75, 10)
+end
+
+@testset "Inverses: spatial_rank=$spatial_rank" for spatial_rank in (1, 2, 3)
+ x = rand(rng, repeat([8], spatial_rank)..., 3, 2)
+ w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
+ cdims = DenseConvDims(x, w; padding=1)
+ y = unfold(x, cdims)
+ z = fold(y, size(x), cdims)
+ divisor = fold(unfold(ones(eltype(x), size(x)...), cdims), size(x), cdims)
+ @test isapprox(z ./ divisor, x, rtol=1.0e-7)
+
+ # introduce stride
+ cdims = DenseConvDims(x, w; padding=1, stride=2)
+ y = unfold(x, cdims)
+ z = fold(y, size(x), cdims)
+ divisor = fold(unfold(ones(eltype(x), size(x)...), cdims), size(x), cdims)
+ @test isapprox(z ./ divisor, x, rtol=1.0e-7)
+end
+
+@testset "AutoDiff: spatial_rank=$spatial_rank" for spatial_rank in (1, 2, 3)
+ x = rand(rng, repeat([5], spatial_rank)..., 3, 2)
+ w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
+ cdims = DenseConvDims(x, w)
+ gradtest(x -> sum(unfold(x, cdims)), x)
+
+ y = unfold(x, cdims)
+ gradtest(y -> sum(fold(y, size(x), cdims)), y)
+end
+

test/runtests.jl

@@ -52,6 +52,10 @@ include("test_utils.jl")
  include("ctc.jl")
  end
 
+ @testset "Fold/Unfold" begin
+ include("fold.jl")
+ end
+
  @testset "Inference" begin
  include("inference.jl")
  end