torch.backends.cudnn.benchmark

2021-04-15 13:04:59 阅读：273 来源： 互联网

标签：args perf benchmark torch algoPerf CUDNN backends ALGO



import torch.backends.cudnn as cudnn
cudnn.benchmark = True

可以在 PyTorch 中对模型里的卷积层进行预先的优化，也就是在每一个卷积层中测试 cuDNN 提供的所有卷积实现算法，然后选择最快的那个。这样在模型启动的时候，只要额外多花一点点预处理时间，就可以较大幅度地减少训练时间。
适用于网络稳定固定输入尺寸大小时候使用，batchsize固定。

// 具体位置的网址：https://github.com/pytorch/pytorch/blob/b5fa9a340a0d174131ad0a452c395860d571b5b0/aten/src/ATen/native/cudnn/Conv.cpp#L701
template<typename perf_t>
void findAlgorithm(const ConvolutionArgs& args, bool benchmark, perf_t* algoPerf) {
  using search = algorithm_search<perf_t>;
  auto& cache = search::cache();

  // 如果缓存里面已经对该卷积场景优化的结果了，那么就直接返回，不找了
  if (cache.find(args.params, algoPerf)) {
    return;
  }

  // 如果在 PyTorch 程序中设置了 torch.backends.cudnn.deterministic=True,
  // 并且 cudnn.benchmark == False的话，那么就选那个默认的卷积算法，返回
  if (args.params.deterministic && !benchmark) {
    algoPerf->algo = search::DEFAULT_ALGO;
    if (args.params.dataType == CUDNN_DATA_HALF) {
      algoPerf->mathType = CUDNN_TENSOR_OP_MATH;
    } else {
      algoPerf->mathType = CUDNN_DEFAULT_MATH;
    }
    search::getWorkspaceSize(args, algoPerf->algo, &(algoPerf->memory));
    return;
  }

  // 再次检查一下缓存中有没有已经对该卷积场景做过选择，
  // recheck 的原因是可能其他线程可能在此期间优化过了
  if (benchmark) {
    if (cache.find(args.params, algoPerf)) {
      // re-check cache since another thread may have benchmarked the algorithm
      return;
    }
  }

  // 好，如果前边三关都过了的话，确实之前没有对该场景做出过优化，
  // 那就调用 search::findAlgorithm 来做 benchmarking。
  // 至于何为 search::findAlgorithm 函数，等等看下边。
  auto perfResults = search::findAlgorithm(args, benchmark);

  // 如果 findAlgorithm 程序运行成功了，并且程序不要求 determinnistic,
  // 使用 findAlgorithm 的结果
  // 否则的话，要求 deterministic，还是返回默认的卷积算法
  // for deterministic algo, look at all the perf results and return the best
  // deterministic algo
  if (perfResults.status == CUDNN_STATUS_SUCCESS &&
      !(args.params.deterministic && perfResults.determinism != CUDNN_DETERMINISTIC)) {

      // if benchmarking, map the original params with the found algo+math type for re-use
      if (benchmark) {
        // cache 只存需要 benchmark 的结果
        cache.insert(args.params, perfResults);

        // Free the cached blocks in our caching allocator. They are
        // needed here because the above benchmarking uses a huge amount of memory,
        // e.g. a few GBs.
        c10::cuda::CUDACachingAllocator::emptyCache();
      }

      *algoPerf = perfResults;
  } else {
      algoPerf->algo = search::DEFAULT_ALGO;
      if (args.params.dataType == CUDNN_DATA_HALF) {
        algoPerf->mathType = CUDNN_TENSOR_OP_MATH;
      } else {
        algoPerf->mathType = CUDNN_DEFAULT_MATH;
      }
      search::getWorkspaceSize(args, algoPerf->algo, &(algoPerf->memory));
  }
}


// 选择卷积 forward 算法的函数
// 具体位置的网址: https://github.com/pytorch/pytorch/blob/b5fa9a340a0d174131ad0a452c395860d571b5b0/aten/src/ATen/native/cudnn/Conv.cpp#L504
template<>
struct algorithm_search<cudnnConvolutionFwdAlgoPerf_t> {
  using perf_t = cudnnConvolutionFwdAlgoPerf_t;
  using algo_t = cudnnConvolutionFwdAlgo_t;

  // 默认算法来了！
  static constexpr auto DEFAULT_ALGO = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM;
  static BenchmarkCache<perf_t>& cache() { return fwd_algos; }

  static perf_t findAlgorithm(const ConvolutionArgs& args, bool benchmark) {
    // CuDNN 实现的 forward 算法，任君选择：
    static const algo_t algos[] = {
         CUDNN_CONVOLUTION_FWD_ALGO_GEMM,
         CUDNN_CONVOLUTION_FWD_ALGO_FFT,
         CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING,
         CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM,
         CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM,
         CUDNN_CONVOLUTION_FWD_ALGO_DIRECT,
         CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD,
         CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED,
    };
    static constexpr int num_algos = CUDNN_CONVOLUTION_FWD_ALGO_COUNT;
    static_assert(sizeof(algos) / sizeof(algos[0]) == num_algos,
                  "Missing cuDNN convolution forward algorithms");
    int perf_count;
    std::unique_ptr<perf_t[]> perf_results(new perf_t[num_algos]);

    // 如果不进行 benchmark 的话，就是我们什么都不设置，PyTorch 默认情况下，
    // 会调用 cudnnGetConvolutionForwardAlgorithm_v7 ！
    if (!benchmark) {
      AT_CUDNN_CHECK(cudnnGetConvolutionForwardAlgorithm_v7(
          args.handle,
          args.idesc.desc(),
          args.wdesc.desc(),
          args.cdesc.desc(),
          args.odesc.desc(),
          num_algos,
          &perf_count,
          perf_results.get()));
    } else { // 如果使用 benchmark，会调用 cudnnFindConvolutionForwardAlgorithmEx ！
      size_t max_ws_size = getMaxWorkspaceSize(args, algos, num_algos);
      Workspace ws(max_ws_size);
      AT_CUDNN_CHECK(cudnnFindConvolutionForwardAlgorithmEx(
          args.handle,
          args.idesc.desc(), args.input.data_ptr(),
          args.wdesc.desc(), args.weight.data_ptr(),
          args.cdesc.desc(),
          args.odesc.desc(), args.output.data_ptr(),
          num_algos,
          &perf_count,
          perf_results.get(),
          ws.data,
          ws.size));
    }
    return getBestAlgorithm<perf_t>(perf_results.get(), args, perf_count);
  }

标签：args,perf,benchmark,torch,algoPerf,CUDNN,backends,ALGO
来源： https://blog.csdn.net/weixin_42917352/article/details/115720974

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

ICode9

torch.backends.cudnn.benchmark