| 112 | |
| 113 | template<class MyDevice> |
| 114 | void Conv2D::forward_dev_impl(const MyDevice & dev, const vector<const Tensor*>& xs, Tensor& fx) const { |
| 115 | DYNET_ASSERT(xs.size() == 2 || xs.size() == 3, "Failed dimension check in Conv2D::forward, at least 2 inputs"); |
| 116 | DYNET_ASSERT(fx.d.bd == xs[0]->d.bd, "Failed dimension check in Conv2D::forward, batchsize not match"); |
| 117 | DYNET_ASSERT(fx.d[2] == xs[1]->d[3], "Failed dimension check in Conv2D::forward, #channel not match"); |
| 118 | AlignedMemoryPool* scratch_allocator = default_device->pools[(int)DeviceMempool::SCS]; |
| 119 | #ifdef __CUDACC__ |
| 120 | #if HAVE_CUDNN |
| 121 | if (!cudnn_conv_op_) |
| 122 | cudnn_conv_op_.reset(new CudnnConvOp(stride, is_valid)); |
| 123 | cudnn_conv_op_->forward_impl(dev, xs, fx); |
| 124 | #else |
| 125 | throw std::runtime_error("Conv2D::forward_dev_impl not supported without CUDNN"); |
| 126 | #endif |
| 127 | #else |
| 128 | Eigen::PaddingType padding_type = is_valid ? Eigen::PADDING_VALID : Eigen::PADDING_SAME; |
| 129 | //void* CHWN_x_mem = aux_mem_pool.allocate(xs[0]->d.size() * sizeof(float)); |
| 130 | void* CHWN_x_mem = scratch_allocator->allocate(xs[0]->d.size() * sizeof(float)); |
| 131 | Tensor CHWN_x = Tensor(Dim({xs[0]->d[2], xs[0]->d[0], xs[0]->d[1]}, xs[0]->d.bd), static_cast<float*>(CHWN_x_mem), xs[0]->device, DeviceMempool::FXS); |
| 132 | Eigen::array<ptrdiff_t, 4> shuffles; |
| 133 | shuffles[0] = 2; shuffles[1] = 0; shuffles[2] = 1; shuffles[3] = 3; |
| 134 | tb<3>(CHWN_x).device(*dev.edevice) = tb<3>(*xs[0]).shuffle(shuffles); |
| 135 | //void* NCHW_f_mem = aux_mem_pool.allocate(xs[1]->d.size() * sizeof(float)); |
| 136 | void* NCHW_f_mem = scratch_allocator->allocate(xs[1]->d.size() * sizeof(float)); |
| 137 | Tensor NCHW_f = Tensor(Dim({xs[1]->d[3], xs[1]->d[2], xs[1]->d[0], xs[1]->d[1]}), static_cast<float*>(NCHW_f_mem), xs[1]->device, DeviceMempool::FXS); |
| 138 | shuffles[0] = 3; shuffles[1] = 2; shuffles[2] = 0; shuffles[3] = 1; |
| 139 | t<4>(NCHW_f).device(*dev.edevice) = t<4>(*xs[1]).shuffle(shuffles); |
| 140 | //void* CHWN_y_mem = aux_mem_pool.allocate(fx.d.size() * sizeof(float)); |
| 141 | void* CHWN_y_mem = scratch_allocator->allocate(fx.d.size() * sizeof(float)); |
| 142 | Tensor CHWN_y = Tensor(Dim({fx.d[2], fx.d[0], fx.d[1]}, fx.d.bd), static_cast<float*>(CHWN_y_mem), fx.device, DeviceMempool::FXS); |
| 143 | tb<3>(CHWN_y).device(*dev.edevice) = Eigen::SpatialConvolution(tb<3>(CHWN_x), t<4>(NCHW_f), stride[0], stride[1], padding_type); |
| 144 | shuffles[0] = 1; shuffles[1] = 2; shuffles[2] = 0; shuffles[3] = 3; |
| 145 | tb<3>(fx).device(*dev.edevice) = tb<3>(CHWN_y).shuffle(shuffles); |
| 146 | if (xs.size() == 3) { |
| 147 | Tensor bias = Tensor(Dim({fx.d[0], fx.d[1], fx.d.bd}, 1), static_cast<float*>(CHWN_x_mem), xs[2]->device, DeviceMempool::FXS); |
| 148 | for (unsigned i = 0; i < fx.d[2]; ++i) { |
| 149 | TensorTools::constant(bias, vec(*xs[2])(i)); |
| 150 | tb<3>(fx).chip<2>(i).device(*dev.edevice) += t<3>(bias); |
| 151 | } |
| 152 | } |
| 153 | #endif |
| 154 | scratch_allocator->free(); |
| 155 | } |
| 156 | |
| 157 | template<class MyDevice> |
| 158 | void Conv2D::backward_dev_impl(const MyDevice & dev, |