Merge #1792

1792: Add trilinear Upsample layer r=ToucheSir a=tknopp

The trilinear Upsample layer was missing although it was actually already implemented in NNLib. This PR makes it accessible through the `Upsample` layer.

### PR Checklist

- [x] Tests are added
- [x] Entry in NEWS.md
- [x] Documentation, if applicable
- [x] API changes require approval from a committer (different from the author, if applicable)


Co-authored-by: Tobias Knopp <tobias@knoppweb.de>

Author: bors[bot]

NEWS.md

@@ -2,6 +2,7 @@
 
 ## v0.12.9
 * Fixed incorrect output and added GPU compatibility for [AlphaDropout](https://github.com/FluxML/Flux.jl/pull/1781).
+* Add trilinear [Upsample layer](https://github.com/FluxML/Flux.jl/pull/1792).
 
 ## v0.12.8
 * Optimized inference and gradient calculation of OneHotMatrix[pr](https://github.com/FluxML/Flux.jl/pull/1756)

docs/src/models/nnlib.md

@@ -56,6 +56,7 @@ NNlib.depthwiseconv
 ```@docs
 NNlib.upsample_nearest
 NNlib.upsample_bilinear
+NNlib.upsample_trilinear
 NNlib.pixel_shuffle
 ```
 

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,))