lib/dnn/core/layers.rb in ruby-dnn-0.1.7 vs lib/dnn/core/layers.rb in ruby-dnn-0.1.8

@@ -186,10 +186,16 @@
         iw2 = iw + pad * 2
         img2 = SFloat.zeros(bs, c, ih2, iw2)
         img2[true, true, pad...(ih + pad), pad...(iw + pad)] = img
         img2
       end
+
+      def back_padding(img, pad)
+        i_end = img.shape[2] - pad
+        j_end = img.shape[3] - pad
+        img[true, true, pad...i_end, pad...j_end]
+      end
     end
     
     
     class Conv2D < HasParamLayer
       include Initializers
@@ -244,11 +250,12 @@
           dridge = @weight_decay * @params[:weight]
           @grads[:weight] += dridge
         end
         @grads[:bias] = dout.sum(0)
         dcol = dout.dot(@params[:weight].transpose)
-        col2im(dcol, @x_shape, @out_height, @out_width, @filter_height, @filter_width, @strides)
+        dx = col2im(dcol, @x_shape, @out_height, @out_width, @filter_height, @filter_width, @strides)
+        @padding ? back_padding(dx, @padding) : dx
       end
     
       def shape
         [@num_filters, @out_height, @out_width]
       end
@@ -291,15 +298,16 @@
     
       def init(model)
         super
         prev_height, prev_width = prev_layer.shape[1], prev_layer.shape[2]
         @num_channel = prev_layer.shape[0]
-        @out_height = (prev_height - @pool_height) / @strides[0] + 1
-        @out_width = (prev_width - @pool_width) / @strides[1] + 1
+        @out_height = (prev_height + @padding * 2 - @pool_height) / @strides[0] + 1
+        @out_width = (prev_width + @padding * 2 - @pool_width) / @strides[1] + 1
       end
     
       def forward(x)
+        x = padding(x, 2) if @padding > 0
         @x_shape = x.shape
         col = im2col(x, @out_height, @out_width, @pool_height, @pool_width, @strides)
         col = col.reshape(x.shape[0] * @out_height * @out_width * x.shape[1], @pool_height * @pool_width)
         @max_index = col.max_index(1)
         col.max(1).reshape(x.shape[0], @out_height, @out_width, x.shape[1]).transpose(0, 3, 1, 2)
@@ -309,10 +317,11 @@
         dout = dout.transpose(0, 2, 3, 1)
         pool_size = @pool_height * @pool_width
         dmax = SFloat.zeros(dout.size * pool_size)
         dmax[@max_index] = dout.flatten
         dcol = dmax.reshape(dout.shape[0..2].reduce(:*), dout.shape[3] * pool_size)
-        col2im(dcol, @x_shape, @out_height, @out_width, @pool_height, @pool_width, @strides)
+        dx = col2im(dcol, @x_shape, @out_height, @out_width, @pool_height, @pool_width, @strides)
+        @padding ? back_padding(dx, @padding) : dx
       end
     
       def shape
         [@num_channel, @out_height, @out_width]
       end