Method Setup

src/model/layer/cudnn_convolution.cc:38–48

}

void CudnnConvolution::Setup(const Shape& in_sample, const LayerConf &conf) {
  Convolution::Setup(in_sample, conf);
  ConvolutionConf conv_conf = conf.convolution_conf();
  // convert MB to bytes
  workspace_byte_limit_ = conv_conf.workspace_byte_limit() << 20;
  prefer_ = ToLowerCase(conv_conf.prefer());
  CHECK(prefer_ == "fastest" || prefer_ == "limited_workspace" ||
        prefer_ == "no_workspace" || prefer_ == "autotune")
      << "CudnnConvolution only supports four algorithm preferences: fastest, "
      "limited_workspace, no_workspace and autotune";
}

void CudnnConvolution::ToDevice(std::shared_ptr<Device> device) {
  Convolution::ToDevice(device);