cleanup

Author: CarloLucibello

src/NNlib.jl

@@ -25,7 +25,6 @@ end
 
 include("activations.jl")
 include("softmax.jl")
-include("misc.jl")
 include("batched/batchedmul.jl")
 include("gemm.jl")
 include("conv.jl")

src/misc.jl

@@ -1,4 +1,4 @@
-export bilinear_upsample, ∇bilinear_upsample
+export bilinear_upsample, ∇bilinear_upsample, pixel_shuffle
 
 """
     bilinear_upsample(x::AbstractArray{<:Number,4}, k::NTuple{2,Int})
@@ -118,13 +118,10 @@ function ∇bilinear_upsample(Δ::AbstractArray{<:Number, 4}, k::NTuple{2,Int})
 
     n_chan, n_batch = size(Δ, 3), size(Δ, 4)
 
-    kern1 = get_downsamplekernel(k[1])
-    kern2 = get_downsamplekernel(k[2])
+    kern1 = get_downsamplekernel(Δ, k[1])
+    kern2 = get_downsamplekernel(Δ, k[2])
     kern = kern1 .* kern2'
-    # TODO there must be a more convenient way to send to gpu 
-    kern = convert(typeof(Δ), reshape(kern, size(kern)..., 1, 1))
-    kern = dropdims(kern, dims=(3,4))
-
+    
     pad = (floor(Int, k[1]//2), floor(Int, k[2]//2))
     stride = k
     weight = similar(Δ, eltype(Δ), (size(kern)..., n_chan, n_chan))
@@ -197,17 +194,9 @@ function ∇bilinear_upsample(Δ::AbstractArray{<:Number, 4}, k::NTuple{2,Int})
     return dx
 end
 
-
-"""
-    get_downsamplekernel(n::Int)
-
-# Arguments
-- `n`: upsample factor for which a downsample kernel will be determined
-
-# Outputs
-- `kernel`: downsample kernel
-"""
-function get_downsamplekernel(n::Int)
+# `n` upsample factor for which a downsample kernel will be determined.
+# Δ is given in case of necessity of gpu conversion 
+function get_downsamplekernel(Δ, n::Int)
     step = 1//n
     if n % 2 == 0
         start = step//2
@@ -218,6 +207,9 @@ function get_downsamplekernel(n::Int)
         upward = collect(start:step:1//1)
         kernel = [upward; reverse(upward[1:end-1])]
     end
+    # TODO there must be a more convenient way to send to gpu 
+    kernel = convert(typeof(Δ), reshape(kernel, length(kernel), 1, 1, 1))
+    kernel = dropdims(kernel, dims=(2,3,4))
     return kernel
 end
 
@@ -228,3 +220,28 @@ function ChainRulesCore.rrule(::typeof(bilinear_upsample), x, k)
     end
     return Ω, bilinear_upsample_pullback
 end
+
+
+"""
+    pixel_shuffle(x, r)
+    
+Pixel shuffling operation. `r` is the upscale factor for shuffling.
+The operation converts an input of size [W,H,r²C,N] to size [rW,rH,C,N]
+Used extensively in super-resolution networks to upsample 
+towards high resolution features.
+
+Reference : https://arxiv.org/pdf/1609.05158.pdf
+"""
+function pixel_shuffle(x::AbstractArray, r::Integer)
+    @assert ndims(x) > 2
+    d = ndims(x) - 2
+    sizein = size(x)[1:d]
+    cin, n = size(x, d+1), size(x, d+2) 
+    @assert cin % r^d == 0
+    cout = cin ÷ r^d
+    # x = reshape(x, sizein..., fill(r, d)..., cout, n) # bug https://github.com/FluxML/Zygote.jl/issues/866
+    x = reshape(x, sizein..., ntuple(i->r, d)..., cout, n)
+    perm = [d+1:2d 1:d]' |> vec  # = [d+1, 1, d+2, 2, ..., 2d, d]
+    x = permutedims(x, (perm..., 2d+1, 2d+2))
+    return reshape(x, ((r .* sizein)..., cout, n))
+end

src/upsample.jl

@@ -26,3 +26,68 @@
         @test Array(g_cpu) ≈ g_cpu  atol=1e-4
     end
 end
+
+@testset "pixel_shuffle" begin
+    x = reshape(1:16, (2, 2, 4, 1))
+    # [:, :, 1, 1] =
+    #     1  3
+    #     2  4
+    # [:, :, 2, 1] =
+    #     5  7
+    #     6  8
+    # [:, :, 3, 1] =
+    #     9  11
+    #     10  12
+    # [:, :, 4, 1] =
+    #     13  15
+    #     14  16
+
+    y_true = [1  9 3 11
+              5 13 7 15
+              2 10 4 12
+              6 14 8 16][:,:,:,:]
+    
+    y = pixel_shuffle(x, 2)
+    @test size(y) == size(y_true)
+    @test y_true == y
+
+    x = reshape(1:32, (2, 2, 8, 1))
+    y_true = zeros(Int, 4, 4, 2, 1)
+    y_true[:,:,1,1] .= [ 1   9  3  11
+                         5  13  7  15
+                         2  10  4  12
+                         6  14  8  16 ]
+
+    y_true[:,:,2,1] .= [ 17  25  19  27
+                         21  29  23  31
+                         18  26  20  28
+                         22  30  24  32]
+
+    y = pixel_shuffle(x, 2)
+    @test size(y) == size(y_true)
+    @test y_true == y
+
+    x = reshape(1:4*3*27*2, (4,3,27,2))
+    y = pixel_shuffle(x, 3)
+    @test size(y) == (12, 9, 3, 2)
+    # batch dimension is preserved 
+    x1 = x[:,:,:,[1]]
+    x2 = x[:,:,:,[2]]
+    y1 = pixel_shuffle(x1, 3)
+    y2 = pixel_shuffle(x2, 3)
+    @test cat(y1, y2, dims=4) == y
+
+    for d in [1, 2, 3]
+        r = rand(1:5)
+        n = rand(1:5)
+        c = rand(1:5)
+        insize = rand(1:5, d)
+        x = rand(insize..., r^d*c, n)
+        
+        y = pixel_shuffle(x, r)
+        @test size(y) == ((r .* insize)..., c, n)
+
+        gradtest(x -> pixel_shuffle(x, r), x)
+    end
+end
+