Merge pull request #690 from AStupidBear/conv1d

fix 1d convolution

Author: maleadt

src/dnn/nnlib.jl

@@ -41,13 +41,25 @@ end
 
 # Convolution
 
+# Since CUDNN does not support 1D convolution, Conv in Flux will give a CUDNNError if the size is 1-dimensional.
+# We have to reshape the CuArray/PoolDims/DenseConvDims to 4D before feeding to CUDNN.
+fix1d(x) = x
+
+fix1d(x::CuArray{T, 3}) where T = reshape(x, size(x, 1), 1, size(x, 2), size(x, 3))
+
+fix1d(cdims::DenseConvDims{1,K,C_in,C_out,S,P,D,F}) where {K,C_in,C_out,S,P,D,F} = 
+  DenseConvDims{2,(K...,1),C_in,C_out,(S...,1),(P...,0,0),(D...,1),F}((cdims.I...,1))
+
+fix1d(pdims::PoolDims{1,K,S,P,D}) where {K,S,P,D,F} = 
+  PoolDims{2,(K...,1),(S...,1),(P...,0,0),(D...,1)}((pdims.I..., 1), pdims.C_in)
+
 function conv!(y::CuArray{T}, x::CuArray{T}, w::CuArray{T}, cdims::DenseConvDims;
                alpha=1, algo=0) where T<:CUDNNFloat
   if version() < v"6"
     all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
   end
-
-  cudnnConvolutionForward(y, x, w, cdims, alpha=alpha, algo=algo)
+  cudnnConvolutionForward(fix1d(y), fix1d(x), fix1d(w), fix1d(cdims), alpha=alpha, algo=algo)
+  return y
 end
 
 function ∇conv_filter!(dw::CuArray{T}, x::CuArray{T}, dy::CuArray{T},
@@ -56,7 +68,8 @@ function ∇conv_filter!(dw::CuArray{T}, x::CuArray{T}, dy::CuArray{T},
     all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
   end
 
-  cudnnConvolutionBackwardFilter(dw, x, dy, cdims, alpha=alpha, algo=algo)
+  cudnnConvolutionBackwardFilter(fix1d(dw), fix1d(x), fix1d(dy), fix1d(cdims), alpha=alpha, algo=algo)
+  return dw
 end
 
 function ∇conv_data!(dx::CuArray{T}, dy::CuArray{T}, w::CuArray{T},
@@ -65,22 +78,23 @@ function ∇conv_data!(dx::CuArray{T}, dy::CuArray{T}, w::CuArray{T},
     all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
   end
 
-  cudnnConvolutionBackwardData(dx, w, dy, cdims, alpha=alpha, algo=algo)
+  cudnnConvolutionBackwardData(fix1d(dx), fix1d(w), fix1d(dy), fix1d(cdims), alpha=alpha, algo=algo)
+  return dx
 end
 
 ∇conv_bias!(db::CuArray{T}, dy::CuArray{T}; alpha=1, beta=0) where T<:CUDNNFloat =
-  cudnnConvolutionBackwardBias(db, dy, alpha=alpha, beta=beta)
+  (cudnnConvolutionBackwardBias(fix1d(db), fix1d(dy), alpha=alpha, beta=beta); return db)
 
 maxpool!(y::CuArray{T}, x::CuArray{T}, pdims::PoolDims) where T<:CUDNNFloat =
-  cudnnPoolingForward(y, x, pdims; mode=0)
+  (cudnnPoolingForward(fix1d(y), fix1d(x), fix1d(pdims); mode=0); return y)
 
 ∇maxpool!(dx::CuArray{T}, dy::CuArray{T}, y::CuArray{T}, x::CuArray{T},
           pdims::PoolDims) where T<:CUDNNFloat =
-  cudnnPoolingBackward(dx, dy, x, y, pdims, mode=0)
+  (cudnnPoolingBackward(fix1d(dx), fix1d(dy), fix1d(x), fix1d(y), fix1d(pdims), mode=0); return dx)
 
 meanpool!(y::CuArray{T}, x::CuArray{T}, pdims::PoolDims) where T<:CUDNNFloat =
-  cudnnPoolingForward(y, x, pdims, mode=1)
+  (cudnnPoolingForward(fix1d(y), fix1d(x), fix1d(pdims), mode=1); return y)
 
 ∇meanpool!(dx::CuArray{T}, dy::CuArray{T}, y::CuArray{T}, x::CuArray{T},
            pdims::PoolDims) where T<:CUDNNFloat =
-  cudnnPoolingBackward(dx, dy, x, y, pdims, mode=1)
+  (cudnnPoolingBackward(fix1d(dx), fix1d(dy), fix1d(x), fix1d(y), fix1d(pdims), mode=1); return dx)

test/dnn.jl

@@ -23,7 +23,7 @@ else
   @test ∇conv_filter(a, c, cdims) ≈ collect(∇conv_filter(da, dc, cdims))
 
   # Test for agreement between CPU NNlib and CuDNN versions, across a variety of kwargs
-  for num_spatial_dims in (2, 3)
+  for num_spatial_dims in (1, 2, 3)
     # Initialize data we'll run our tests over
     C_in = 3
     C_out = 4