add pixel_shuffle

.github/workflows/ci.yml

@@ -48,13 +48,13 @@ jobs:
     steps:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1
-        # `allow-failure` not available yet https://github.com/actions/toolkit/issues/399
+        ## `allow-failure` not available yet https://github.com/actions/toolkit/issues/399
         continue-on-error: ${{ matrix.julia-version == 'nightly' }}
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
       - uses: actions/cache@v1
-        continue-on-error: ${{ matrix.julia-version == 'nightly' }}
+        # continue-on-error: ${{ matrix.julia-version == 'nightly' }}
         env:
           cache-name: cache-artifacts
         with:
@@ -65,11 +65,11 @@ jobs:
             ${{ runner.os }}-test-
             ${{ runner.os }}-
       - uses: julia-actions/julia-buildpkg@v1
-        continue-on-error: ${{ matrix.julia-version == 'nightly' }}
+        # continue-on-error: ${{ matrix.julia-version == 'nightly' }}
       - uses: julia-actions/julia-runtest@v1
-        continue-on-error: ${{ matrix.julia-version == 'nightly' }}
+        # continue-on-error: ${{ matrix.julia-version == 'nightly' }}
       - uses: codecov/codecov-action@v1
-        continue-on-error: ${{ matrix.version == 'nightly' }}
+        # continue-on-error: ${{ matrix.version == 'nightly' }}
         with:
           file: lcov.info
 

src/NNlib.jl

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

src/misc.jl

@@ -0,0 +1,22 @@
+export pixel_shuffle
+
+"""
+    pixel_shuffle(x, r)
+
+Pixel shuffling operation. `r` is the scale 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{T, 4}, r::Integer) where T <:Number
+    w, h, c, n = size(x)
+    @assert c % r^2 == 0
+    c_out = c ÷ r^2
+    w_out = w * r
+    h_out = h * r
+    x = reshape(x, (w, h, r, r, c_out, n))
+    x = permutedims(x, (3, 1, 4, 2, 5, 6))
+    return reshape(x, (w_out, h_out, c_out, n))
+end

test/misc.jl

@@ -0,0 +1,34 @@
+@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) == (4, 4, 1, 1)
+    @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)
+    x1 = x[:,:,:,[1]]
+    x2 = x[:,:,:,[2]]
+    y1 = pixel_shuffle(x1, 3)
+    y2 = pixel_shuffle(x2, 3)
+    @test cat(y1, y2, dims=4) == y
+    gradtest(x -> pixel_shuffle(x, 2), rand(3,3,8,2))
+end

test/runtests.jl

@@ -35,3 +35,7 @@ end
 @testset "Softmax" begin
     include("softmax.jl")
 end
+
+@testset "Misc Stuff" begin
+    include("misc.jl")
+end