add trilinear Upsample layer

Author: tknopp

src/layers/upsample.jl

@@ -12,6 +12,7 @@ Currently supported upsampling `mode`s
 and corresponding NNlib's methods are:
   - `:nearest` -> [`NNlib.upsample_nearest`](@ref) 
   - `:bilinear` -> [`NNlib.upsample_bilinear`](@ref)
+  - `:trilinear` -> [`NNlib.upsample_trilinear`](@ref)
 
 # Examples
 
@@ -35,7 +36,7 @@ struct Upsample{mode, S, T}
 end
 
 function Upsample(mode::Symbol = :nearest; scale = nothing, size = nothing)
-  mode in [:nearest, :bilinear] || 
+  mode in [:nearest, :bilinear, :trilinear] || 
     throw(ArgumentError("mode=:$mode is not supported."))
   if !(isnothing(scale) ⊻ isnothing(size))
     throw(ArgumentError("Either scale or size should be specified (but not both)."))
@@ -58,6 +59,11 @@ end
 (m::Upsample{:bilinear, Nothing})(x::AbstractArray) = 
   NNlib.upsample_bilinear(x; size=m.size)
 
+(m::Upsample{:trilinear})(x::AbstractArray) =
+  NNlib.upsample_trilinear(x, m.scale)
+(m::Upsample{:trilinear, Nothing})(x::AbstractArray) = 
+  NNlib.upsample_trilinear(x; size=m.size)
+
 function Base.show(io::IO, u::Upsample{mode}) where {mode}
   print(io, "Upsample(")
   print(io, ":", mode)

test/layers/upsample.jl

@@ -18,6 +18,26 @@
   @test size(y) == (4, 6, 2, 3)
 end
 
+@testset "upsample trilinear" begin
+  m = Upsample(:trilinear, scale=(2, 3, 2))
+  x = rand(Float32, 3, 4, 2, 3, 4)
+  y = m(x)
+  @test y isa Array{Float32, 5} 
+  @test size(y) == (6, 12, 4, 3, 4)
+
+  m = Upsample(:trilinear, scale=3)
+  x = rand(Float32, 3, 4, 2, 3, 4)
+  y = m(x)
+  @test y isa Array{Float32, 5} 
+  @test size(y) == (9, 12, 6, 3, 4)
+
+  m = Upsample(:trilinear, size=(4, 6, 4))
+  x = rand(Float32, 3, 4, 2, 3, 4)
+  y = m(x)
+  @test y isa Array{Float32, 5} 
+  @test size(y) == (4, 6, 4, 3, 4)
+end
+
 @testset "upsample nearest" begin
   x = rand(Float32, 3, 2, 3)
   m = Upsample(:nearest, scale=(2,))