Add support for 3D convolutional backbones

src/Vision/models/xresnet.jl

@@ -1,70 +1,107 @@
 
 const act_fn = relu
 
-convx(ni, nf; ks=3, stride=1) = Conv(
-        (ks, ks), ni => nf, stride = stride, pad = ks ÷ 2, init = Flux.kaiming_normal)
+convx(ni, nf; ks = 3, stride = 1, ndim = 2) = Conv(
+    ntuple(_ -> ks, ndim),
+    ni => nf,
+    stride = stride,
+    pad = ks ÷ 2,
+    init = Flux.kaiming_normal,
+)
 
-function convxlayer(ni, nf; ks=3, stride=1, zero_bn=false, act=true)
+function convxlayer(ni, nf; ks = 3, stride = 1, zero_bn = false, act = true, ndim = 2)
     bn = BatchNorm(nf, act ? act_fn : identity)
     fill!(bn.γ, zero_bn ? 0 : 1)
-    return Chain(convx(ni, nf; ks = ks, stride = stride), bn)
+    return Chain(convx(ni, nf; ks = ks, stride = stride, ndim = ndim), bn)
 end
 
 struct ResBlock
-    convs
-    idconv
-    pool
+    convs::Any
+    idconv::Any
+    pool::Any
 end
 
-function ResBlock(expansion::Int, ni::Int, nh::Int; stride::Int = 1)
+function ResBlock(expansion::Int, ni::Int, nh::Int; stride::Int = 1, ndim = 2)
     nf, ni = nh * expansion, ni * expansion
     if expansion == 1
         layers = [
-            convxlayer(ni, nh; stride = stride),
-            convxlayer(nh, nf; zero_bn = true, act = false),
+            convxlayer(ni, nh; stride = stride, ndim = ndim),
+            convxlayer(nh, nf; zero_bn = true, act = false, ndim = ndim),
         ]
     else
         layers = [
-            convxlayer(ni, nh; ks = 1),
-            convxlayer(nh, nh; stride = stride),
-            convxlayer(nh, nf; ks = 1, zero_bn = true, act = false)
+            convxlayer(ni, nh; ks = 1, ndim = ndim),
+            convxlayer(nh, nh; stride = stride, ndim = ndim),
+            convxlayer(nh, nf; ks = 1, zero_bn = true, act = false, ndim = ndim),
         ]
     end
 
     return ResBlock(
         Chain(layers...),
-        ni == nf ? identity : convxlayer(ni, nf; ks = 1, stride = 1),
-        stride == 1 ? identity : MeanPool((2, 2))
+        ni == nf ? identity : convxlayer(ni, nf; ks = 1, stride = 1, ndim = ndim),
+        stride == 1 ? identity : MeanPool(ntuple(_ -> 2, ndim)),
     )
 end
 Flux.@functor ResBlock
 
 
 (r::ResBlock)(x) = act_fn.(r.convs(x) .+ r.idconv(r.pool(x)))
 
-function make_layer(expansion, ni, nf, n_blocks, stride)
-    return Chain([
-            ResBlock(expansion, i == 1 ? ni : nf, nf; stride =(i == 1 ? stride : 1))
-            for i in 1:n_blocks
-            ]...)
+function make_layer(expansion, ni, nf, n_blocks, stride; ndim = 2)
+    return Chain(
+        [
+            ResBlock(
+                expansion,
+                i == 1 ? ni : nf,
+                nf;
+                stride = (i == 1 ? stride : 1),
+                ndim = ndim,
+            ) for i = 1:n_blocks
+        ]...,
+    )
 end
 
-function XResNet(expansion, layers; c_in = 3)
-    nfs = [c_in, (c_in+1)*8, 64, 64]
-    stem = [convxlayer(nfs[i], nfs[i+1]; stride = i == 1 ? 2 : 1) for i in 1:3]
+"""
+    XResNet(expansion, layers; c_in = 3, ndim = 2)
+
+Create an XResNet model backbone following the [implementation in
+fastai](https://github.com/fastai/fastai/blob/master/fastai/vision/models/xresnet.py#L22).
+
+- `c_in::Int = 3`: The number of input channels, e.g. `1` for grayscale images and `3` for
+    RGB images
+- `ndim::Int = 2`: The number of dimensions for the convolutional and pooling layers, e.g.
+    `2` for 2D input images and `3` for 3D volumes.
+"""
+function XResNet(expansion, layers; c_in = 3, ndim = 2)
+    nfs = [c_in, (c_in + 1) * 8, 64, 64]
+    stem = [convxlayer(nfs[i], nfs[i+1]; stride = i == 1 ? 2 : 1, ndim = ndim) for i = 1:3]
 
     nfs = [64 ÷ expansion, 64, 128, 256, 512]
     res_layers = [
-        make_layer(expansion, nfs[i], nfs[i+1], l, i == 1 ? 1 : 2)
-        for (i, l) in enumerate(layers)
+        make_layer(expansion, nfs[i], nfs[i+1], l, i == 1 ? 1 : 2, ndim = ndim) for
+        (i, l) in enumerate(layers)
     ]
 
-    return Chain(
-        stem...,
-        MaxPool((3, 3); pad = 1, stride = 2),
-        res_layers...,
-    )
+    return Chain(stem..., MaxPool(ntuple(_ -> 3, ndim); pad = 1, stride = 2), res_layers...)
 end
 
 xresnet18(; kwargs...) = XResNet(1, [2, 2, 2, 2]; kwargs...)
 xresnet50(; kwargs...) = XResNet(4, [3, 4, 6, 3]; kwargs...)
+
+
+
+@testset "XResNet [model]" begin
+    @testset "Basic" begin
+        @test_nowarn begin
+            model = xresnet18()
+            @test Flux.outputsize(model, (128, 128, 3, 1)) == (4, 4, 512, 1)
+        end
+    end
+
+    @testset "3D" begin
+        @test_nowarn begin
+            model = xresnet18(ndim=3)
+            @test Flux.outputsize(model, (128, 128, 128, 3, 1)) == (4, 4, 4, 512, 1)
+        end
+    end
+end