Flip name to transpose_switchbatch()

Also specialize on kernel size, as that turns out to be helpful for performance.

Author: staticfloat

src/dim_helpers.jl

@@ -6,16 +6,16 @@ include("dim_helpers/PoolDims.jl")
 
 
 """
-    transpose_flipbatch(x::AbstractArray)
+    transpose_swapbatch(x::AbstractArray)
 
-Given an AbstractArray, flip its batch and channel axes, as we must during transposed
+Given an AbstractArray, swap its batch and channel axes, as we must during transposed
 convolution.  We do this to the operands during convolution, and then again to the
 output once we're done.
 """
-function transpose_flipbatch(x::AbstractArray)
+function transpose_swapbatch(x::AbstractArray)
     return permutedims(x, ((1:(ndims(x)-2))..., ndims(x), ndims(x)-1))
 end
-function transpose_flipbatch(x::Tuple)
+function transpose_swapbatch(x::Tuple)
     return (x[1:end-2]..., x[end], x[end-1])
 end
 

src/dim_helpers/DenseConvDims.jl

@@ -5,18 +5,15 @@ export DenseConvDims
 
 Concrete subclass of `ConvDims` for a normal, dense, conv2d/conv3d.
 """
-struct DenseConvDims{N,S,P,D,F} <: ConvDims{N,S,P,D,F}
+struct DenseConvDims{N,K,C_in,C_out,S,P,D,F} <: ConvDims{N,S,P,D,F}
     I::NTuple{N,Int}
-    K::NTuple{N,Int}
-    C_in::Int
-    C_out::Int
 end
 
 # Getters for the fields
 input_size(c::DenseConvDims) = c.I
-kernel_size(c::DenseConvDims) = c.K
-channels_in(c::DenseConvDims) = c.C_in
-channels_out(c::DenseConvDims) = c.C_out
+kernel_size(c::DenseConvDims{N,K,C_in,C_out,S,P,D,F}) where {N,K,C_in,C_out,S,P,D,F} = K
+channels_in(c::DenseConvDims{N,K,C_in,C_out,S,P,D,F}) where {N,K,C_in,C_out,S,P,D,F} = C_in
+channels_out(c::DenseConvDims{N,K,C_in,C_out,S,P,D,F}) where {N,K,C_in,C_out,S,P,D,F} = C_out
 
 # Convenience wrapper to create DenseConvDims objects
 function DenseConvDims(x_size::NTuple{M}, w_size::NTuple{M};
@@ -34,22 +31,16 @@ function DenseConvDims(x_size::NTuple{M}, w_size::NTuple{M};
     # The type parameters are what 
     return DenseConvDims{
         M - 2,
+        w_size[1:end-2],
+        x_size[end-1],
+        w_size[end],
         stride,
         padding,
         dilation,
         flipkernel
     }(
-        # Image spatial size
+        # Input spatial size
         x_size[1:end-2],
-
-        # Kernel spatial size
-        w_size[1:end-2],
-
-        # Input channels
-        x_size[end-1],
-
-        # Output channels
-        w_size[end],
     )
 end
 
@@ -66,7 +57,7 @@ end
 function DenseConvDims(c::ConvDims; N=spatial_dims(c), I=input_size(c), K=kernel_size(c),
                        C_in=channels_in(c), C_out=channels_out(c), S=stride(c),
                        P=padding(c), D=dilation(c), F=flipkernel(c))
-    return DenseConvDims{N, S, P, D, F}(I, K, C_in, C_out)
+    return DenseConvDims{N, K, C_in, C_out, S, P, D, F}(I)
 end
 
 function check_dims(x::NTuple{M}, w::NTuple{M}, y::NTuple{M}, cdims::DenseConvDims) where {M}

src/dim_helpers/PoolDims.jl

@@ -13,11 +13,7 @@ struct PoolDims{N,K,S,P,D} <: ConvDims{N, S, P, D, false}
 end
 
 # Getters for both type parameters and fields
-spatial_dims(c::PoolDims{N,K,S,P,D}) where {N, K, S, P, D} = N
 kernel_size(c::PoolDims{N,K,S,P,D}) where {N, K, S, P, D} = K
-stride(c::PoolDims{N,K,S,P,D}) where {N, K, S, P, D} = S
-padding(c::PoolDims{N,K,S,P,D}) where {N, K, S, P, D} = P
-dilation(c::PoolDims{N,K,S,P,D}) where {N, K, S, P, D} = D
 input_size(c::PoolDims) = c.I
 channels_in(c::PoolDims) = c.C_in
 channels_out(c::PoolDims) = c.C_in

src/impl/conv_im2col.jl

@@ -163,7 +163,12 @@ out along the rows of `col`, one for each output pixel.  This routine is used by
 im2col-based convolutions, just with extra singleton dimensions added in the case of `2d`
 or `1d` images.
 """
-function im2col!(col::AbstractArray{T,2}, x::AbstractArray{T,4}, cdims::ConvDims) where T
+function im2col!(col::AbstractArray{T,2}, x::AbstractArray{T,4},
+                                          cdims::ConvDims) where {T}
+    if spatial_dims(cdims) != 3
+        throw(DimensionMismatch("im2col!() only accepts 3d convoluitional inputs"))
+    end
+
     # Extract those nice, compile-time constant type parameters from `cdims`.
     width, height, depth = input_size(cdims)
     kernel_w, kernel_h, kernel_d = kernel_size(cdims)