add 'size' for each layer type

Author: schillic

src/Architecture/Architecture.jl

@@ -9,6 +9,7 @@ using Requires
 using LinearAlgebra: dot
 using Statistics: mean
 
+import Base: size
 export AbstractNeuralNetwork, FeedforwardNetwork,
        AbstractLayerOp, DenseLayerOp, ConvolutionalLayerOp, FlattenLayerOp,
        AbstractPoolingLayerOp, MaxPoolingLayerOp, MeanPoolingLayerOp,

src/Architecture/LayerOps/ConvolutionalLayerOp.jl

@@ -115,6 +115,8 @@ function Base.show(io::IO, L::ConvolutionalLayerOp)
     return print(io, str)
 end
 
+size(::ConvolutionalLayerOp) = (3, 3)
+
 function load_Flux_convert_Conv_layer()
     return quote
         function Base.convert(::Type{ConvolutionalLayerOp}, layer::Flux.Conv)

src/Architecture/LayerOps/DenseLayerOp.jl

@@ -73,6 +73,8 @@ dim_in(L::DenseLayerOp) = size(L.weights, 2)
 
 dim_out(L::DenseLayerOp) = length(L.bias)
 
+size(::DenseLayerOp) = (1, 1)
+
 function load_Flux_convert_Dense_layer()
     return quote
         function Base.convert(::Type{DenseLayerOp}, layer::Flux.Dense)

src/Architecture/LayerOps/FlattenLayerOp.jl

@@ -42,3 +42,5 @@ function (L::FlattenLayerOp)(T)
     end
     return vec(permutedims(T, (2, 1, 3:length(s)...)))
 end
+
+size(::FlattenLayerOp) = (nothing, 1)

src/Architecture/LayerOps/PoolingLayerOp.jl

@@ -90,3 +90,5 @@ function (L::AbstractPoolingLayerOp)(T)
     end
     return O
 end
+
+size(::AbstractPoolingLayerOp) = (3, 3)

src/Architecture/NeuralNetworks/FeedforwardNetwork.jl

@@ -32,14 +32,20 @@ end
 function _first_inconsistent_layer(L)
     prev = nothing
     for (i, l) in enumerate(L)
-        if !isnothing(prev) && dim_in(l) != dim_out(prev)
+        if !isnothing(prev) &&
+           ((!isnothing(dim_in(l)) && !isnothing(dim_out(prev)) && dim_in(l) != dim_out(prev)) ||
+            !_iscompatible(size(prev), size(l)))
             return i
         end
         prev = l
     end
     return 0
 end
 
+_iscompatible(t1::Tuple, t2::Tuple) = _iscompatible(t1[2], t2[1])
+_iscompatible(i::Int, j::Int) = i == j
+_iscompatible(i, ::Nothing) = true
+
 layers(N::FeedforwardNetwork) = N.layers
 
 function load_Flux_convert_network()

test/Architecture/ConvolutionalLayerOp.jl

@@ -54,6 +54,9 @@ println(io, L)
       !(L ≈ ConvolutionalLayerOp(Ws, [b1 .+ 1], ReLU())) &&
       !(L ≈ ConvolutionalLayerOp(Ws, bs, Id()))
 
+# size
+@test size(L) == (3, 3)
+
 # kernel size and number of filters
 @test kernel(L) == kernel(L2) == (2, 2, 1)
 @test n_filters(L) == 1 && n_filters(L2) == 2

test/Architecture/DenseLayerOp.jl

@@ -42,6 +42,7 @@ println(io, L)
 @test dim_out(L) == 3
 @test dim(L) == (2, 3)
 @test length(L) == 3
+@test size(L) == (1, 1)
 
 # test methods for all activations
 function test_layer(L::DenseLayerOp{Id})

test/Architecture/FeedforwardNetwork.jl

@@ -48,3 +48,16 @@ println(io, N2)
 @test dim_in(N1) == 2 && dim_in(N2) == 2
 @test dim_out(N1) == 3 && dim_out(N2) == 2
 @test dim(N1) == (2, 3) && dim(N2) == (2, 2)
+
+# network with all layer types
+L1 = ConvolutionalLayerOp([reshape([1 0; -1 2], (2, 2, 1))], [1], ReLU())
+L2 = MaxPoolingLayerOp(1, 1)
+L3 = FlattenLayerOp()
+W = zeros(2, 9); W[1, 1] = W[2, 2] = 1
+L4 = DenseLayerOp(W, [1.0 0], ReLU())
+N3 = FeedforwardNetwork([L1, L2, L3, L4])
+T441 = reshape([0 4 2 1; -1 0 1 -2; 3 1 2 0; 0 1 4 1], (4, 4, 1))
+@test N3(T441) == [3.0 2; 8 7]
+
+# incompatible dimensions
+@test_throws ArgumentError FeedforwardNetwork([L1, L4])

test/Architecture/FlattenLayerOp.jl

@@ -6,3 +6,6 @@ L = FlattenLayerOp()
 
 # equality
 @test L == FlattenLayerOp()
+
+# size
+@test size(L) == (nothing, 1)

test/Architecture/PoolingLayerOp.jl

@@ -39,3 +39,6 @@ for (L, LT) in zip(Ls, LTs)
     @test L != LT(2, 2) && L != LT(3, 3)
 end
 @test L1 != L2
+
+# size
+@test size(L1) == (3, 3)