背景知识
首先,要了解benchmark的作用,我们需要了解cuDNN。cuDNN是英伟达为深度神经网络开发的GPU加速工具库,针对卷积、池化等操作做了底层优化。Pytorch、Tensorflow等都默认使用cuDNN做加速操作。
卷积操作
例如对图像进行采样,通常的做法是使用多层循环嵌套,对每一个通道的每一个区域,用同样的卷积核进行卷积操作,然后继续滑动,直到direct 到整张图像。这个方法需要多次迭代,效率较低。除此之外,还可以使用 GEMM (General Matrix Multiply) ,FFT(Fast Fourier Transform),Winograd等算法。通常情况下,在使用大的卷积核是会具有更快的速度,而且每层算法的卷积核大小是不同(这种优化更符合实际的情况)。卷积核的大小并不影响卷积的时间,但图像的尺寸大小会影响卷积的效率。
例如:
我们只要固定输入大小都是 (8, 64, 224, 224),即 batch_size 为 8,输入的通道为 64,宽和高为 224,那么卷积层的运行时间都是几乎不变的,无论其中每个像素具体的值是 0.1 还是 1000.0。
这样固定了模型输入的尺寸大小,所以对每个卷积层来说,其接受的输入尺寸都是静态的。
torch.backends.cudnn.benchmark
Pytorch中的backends是在调用底层库。其中包括:
cuda
cudnn
mkl
mkldnn
openmp
当我们将flag设置为True时,cuDNN会对每一个卷积层进行测试,找到最优的算法。这样就使得在开始训练时的时间会增加,随着训练的继续,训练时间将减小。但在一个变化的网络模型中使用backends,每轮次训练都需要计算最优,反而增加了耗时。
Pytorch中的源码:
// 具体位置的网址: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);
}
参考:
torch.backends.cudnn.benchmark ?!