single-batch convolution

src/layers/conv.jl

@@ -33,6 +33,7 @@ Conv(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity; init =
 Flux.treelike(Conv)
 
 function (c::Conv)(x)
+  ndims(x) == ndims(c.weight)-1 && return squeezebatch(c(reshape(x, size(x)..., 1)))
   σ, b = c.σ, reshape(c.bias, map(_->1, c.stride)..., :, 1)
   σ.(conv(x, c.weight, stride = c.stride, pad = c.pad) .+ b)
 end

src/utils.jl

@@ -46,6 +46,10 @@ function frequencies(xs)
   return fs
 end
 
+head(x::Tuple) = reverse(Base.tail(reverse(x)))
+
+squeezebatch(x) = reshape(x, head(size(x)))
+
 """
   batch(xs)
 

test/tracker.jl

@@ -92,4 +92,6 @@ end
 
 @test @sprintf("%.2f", sum(param([1,2,3]))) == "6.00"
 
+@inferred NNlib.conv(param(rand(10,10,3,2)),randn(2,2,3,4))
+
 end #testset