params Class — pytorch Architecture

Source Code

aten/src/ATen/native/cudnn/Conv_v7.cpp lines 108–132

template <typename T>
struct BenchmarkCache {
  std::mutex mutex;
  std::unordered_map<
      ConvolutionParams,
      T,
      ParamsHash<ConvolutionParams>,
      ParamsEqual<ConvolutionParams>>
      map;

  bool find(const ConvolutionParams& params, T* results) {
    std::lock_guard<std::mutex> guard(mutex);
    auto it = map.find(params);
    if (it == map.end()) {
      return false;
    }
    *results = it->second;
    return true;
  }

  void insert(const ConvolutionParams& params, const T& results) {
    std::lock_guard<std::mutex> guard(mutex);
    map[params] = results;
  }
};