Notas de estudio de código MMDetection3D-torch.backends.cudnn.benchmark
Prefacio
Cuando entrena modelos de aprendizaje profundo, a menudo usa GPU para acelerar el entrenamiento de la red. Pero cuando se trata de torch.backends.cudnn.benchmark, la bandera relacionada con la GPU, algunas personas pueden sentirse desconocidas. En escenarios generales, siempre que el valor se establezca en Verdadero al comienzo del programa PyTorch, la velocidad de ejecución de la red neuronal convolucional se puede mejorar considerablemente. Dado que es tan sorprendente, ¿por qué PyTorch no lo establece en True de forma predeterminada? ¿Cuáles son sus escenarios aplicables? ¿Por qué utilizarlo para mejorar la eficiencia? La respuesta está en este artículo.
Introducción a la versión condensada
Establecer torch.backends.cudnn.benchmark = True hará que el programa dedique un poco de tiempo extra al principio para buscar el algoritmo de implementación de convolución más adecuado para cada capa convolucional de toda la red, logrando así la aceleración de la red. El escenario aplicable es que la estructura de la red es fija (no cambia dinámicamente) y la forma de entrada de la red (incluido el tamaño del lote, el tamaño de la imagen y el canal de entrada) no cambia. De hecho, es más aplicable en circunstancias normales. Por el contrario, si la configuración de la capa convolucional sigue cambiando, hará que el programa se optimice continuamente, lo que consumirá más tiempo.
Introducción detallada a la versión
conocimiento de fondo
Antes de hablar sobre torch.backends.cudnn.benchmark, primero presentamos brevemente cuDNN. cuDNN es una biblioteca de aceleración de GPU especialmente desarrollada por NVIDIA para redes neuronales profundas. Ha realizado muchas optimizaciones de bajo nivel para operaciones comunes como convolución y agrupación, y es mucho más rápido que los programas de GPU generales. La mayoría de los marcos de aprendizaje profundo convencionales admiten cuDNN, y PyTorch no es una excepción. Al usar GPU, PyTorch usará la aceleración cuDNN de forma predeterminada. Sin embargo, cuando se usa cuDNN, el modo torch.backends.cudnn.benchmark es Falso. ¡Eso significa que nuestro programa puede continuar acelerándose!
La capa convolucional es la parte más importante de la red neuronal convolucional y, a menudo, es la parte más computacionalmente intensiva. Si podemos mejorar la eficiencia de las operaciones de convolución en el código subyacente, podemos aumentar en gran medida la velocidad de entrenamiento y predicción sin cambiar la estructura de la red neuronal dada.
Para el funcionamiento de la convolución, hay muchas formas de implementarlo. La forma más sencilla de lograrlo es utilizar el anidamiento de bucle multicapa. Para cada imagen de entrada, para cada canal a emitir, para cada canal de entrada, seleccione un área, realice la convolución con el kernel de convolución especificado y luego proceda uno por uno .Deslice hasta que se procese toda la imagen. Este método generalmente se llama método directo. Aunque este método es simple, sabemos que la eficiencia no será muy alta cuando veamos tantos ciclos. Además, existen algoritmos para implementar capas convolucionales basados en GEMM (General Matrix Multiply), basado en FFT, algoritmo de Winograd, etc., y cada algoritmo tiene sus propias variantes. En una biblioteca triNNity de código abierto C ++, ¡se han implementado casi 80 algoritmos de propagación hacia adelante para convolución!
Cada algoritmo de convolución tiene algunas ventajas únicas. Por ejemplo, algunos algoritmos son muy rápidos cuando el núcleo de convolución es grande; por ejemplo, algunos algoritmos usan menos memoria en algunos casos. Dada una red neuronal convolucional (como ResNet-101), dado el tamaño de la imagen de entrada, dada la plataforma de hardware, la forma más fácil de implementar esta red es utilizar el mismo algoritmo de convolución para todas las capas convolucionales (como el algoritmo directo )), pero esta operación definitivamente no es óptima; una mejor manera es que podemos realizar algunas pruebas de optimización simples de antemano, seleccionar el algoritmo de convolución más adecuado (más rápido) en cada capa de convolución y decidir cada uno después del algoritmo más rápido en la capa , luego ejecutamos toda la red, lo que mejorará mucho la eficiencia.
Aquí hay una pregunta, ¿por qué podemos elegir el algoritmo para cada capa de antemano, incluso si las imágenes que enviamos a la red para el entrenamiento son diferentes cada vez? Es decir, la entrada de la red cambia cada vez, entonces, ¿cómo puedo asegurarme de que el algoritmo óptimo seleccionado de antemano también sea adecuado para esta entrada? La razón es que para una entrada dada, el tamaño de su valor específico no afecta el tiempo de ejecución de la convolución, solo afectará su tamaño. Por ejemplo, siempre que el tamaño de entrada fijo sea (8, 64, 224, 224), es decir, el tamaño de lote es 8, el canal de entrada es 64 y el ancho y alto son 224, entonces el tiempo de ejecución del convolucional La capa casi no cambia, no importa que el valor específico de cada píxel sea 0,1 o 1000,0.
En este caso, debido a que fijamos el tamaño de la entrada del modelo, para cada capa convolucional, el tamaño de entrada aceptado es estático y fijo. Solo necesitamos usar la inicialización aleatoria cuando optimizamos por adelantado. La entrada del tamaño correspondiente se puede probar y seleccionado.
antorcha.backends.cudnn.benchmark！
Habiendo dicho tanto conocimiento previo, ¿cuál es la conexión con cudnn.benchmark? De hecho, al establecer esta bandera en Verdadero, podemos optimizar previamente la capa convolucional en el modelo en PyTorch, es decir, probar todos los algoritmos de implementación de convolución proporcionados por cuDNN en cada capa convolucional, y luego elegir el más rápido. De esta forma, cuando se inicia el modelo, siempre que se invierta un poco más de tiempo de preprocesamiento, el tiempo de entrenamiento se puede reducir considerablemente.
¿No es así, es genial usar cudnn.benchmark por un tiempo, pero siempre es genial usarlo todo el tiempo? De hecho, de lo contrario, en algunos casos, su uso puede aumentar considerablemente el tiempo de ejecución. Ya lo hemos mencionado en el conocimiento previo, pero aquí lo definimos más específicamente, qué factores afectan el tiempo de ejecución de la capa convolucional.
1.首先，当然是卷积层本身的参数，常见的包括卷积核大小，stride，dilation，padding ，输出通道的个数等；
2.其次，是输入的相关参数，包括输入的宽和高，输入通道的个数等；
3.最后，还有一些其他的因素，比如硬件平台，输入输出精度、布局等等。
Los parámetros que definimos en una escena convolucional incluyen principalmente (1) y (2), porque (3) suelen ser los mismos en el mismo programa, por lo que los ignoramos por el momento. Los diferentes escenarios de convolución tienen diferentes algoritmos de convolución óptimos, que deben probarse y seleccionarse por separado.
Basándonos en esto, podemos ver que, en primer lugar, si nuestro modelo de red sigue cambiando, no debe establecerse cudnn.benchmark = True. Debido a que la estructura de la red cambia a menudo, PyTorch la optimizará automáticamente de acuerdo con la nueva escena de convolución cada vez: esta vez tomó mucho tiempo seleccionar el algoritmo más adecuado. Como resultado, la próxima vez que cambie su estructura, la optimizará para nada. No solo eso, sino que tenemos que seguir eligiendo la combinación de algoritmos más eficiente basada en esta nueva estructura, lo que lleva mucho tiempo. Esto reducirá en gran medida la eficiencia.
Además, el tamaño de nuestra entrada no se puede cambiar. Para una capa convolucional, por ejemplo, la forma de entrada esta vez es (8, 3, 224, 224), y la próxima vez que se cambie a (8, 3, 112, 112), eso no funcionará. La situación de entrada ha cambiado y es posible que el algoritmo óptimo no sea aplicable (por ejemplo, algunos algoritmos son rápidos en el caso de entrada de gran tamaño), y PyTorch seguirá encontrando el algoritmo óptimo nuevamente. Tenga en cuenta que el tamaño del lote, el canal de entrada y el tamaño de la imagen aquí no se pueden cambiar.
Pero para el modelo de CV general, la estructura de la red generalmente no cambia dinámicamente. En segundo lugar, la imagen generalmente cambia de tamaño a un tamaño fijo y el tamaño del lote también es fijo. Por lo tanto, en la mayoría de los casos, podemos agregar esta línea de código mágico al programa para reducir el tiempo de ejecución.
¿Dónde debería agregarse esta línea de código?
He visto algunos blogs en Internet que mencionan el uso de cudnn.benchmark = True, pero no especificaron dónde colocar este código. Me temo que algunos lectores aquí son los mismos que yo en ese momento, así que agregué esta pequeña parte. De hecho, generalmente es bueno agregarlo al principio. Por ejemplo, cuando configure el uso de GPU, agregue una oración después:
if args.use_gpu and torch.cuda.is_available():
device = torch.device('cuda')
torch.backends.cudnn.benchmark = True
else:
device = torch.device('cpu')
......
......
Por supuesto, en algunos casos, también puede cambiar el valor de torch.backends.cudnn.benchmark varias veces en el programa para hacer trucos o algo así.
Código fuente correspondiente en PyTorch
Todo esto es de lo que estaba hablando, así que echemos un vistazo al código fuente en PyTorch y veamos su lógica original. El código proviene de. Esta parte es el núcleo de la selección del algoritmo de convolución en cuDNN. Agregué brevemente algunas oraciones Notas:
// 具体位置的网址：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);
}
Bien, he terminado de leer los dos códigos fuente principales. Puede que tengas dos dudas aquí. ¿Cuál es el primer torch.backends.cudnn.deterministic? Como su nombre lo indica, si esta bandera se establece en Verdadero, se determinará el algoritmo de convolución devuelto cada vez, es decir, el algoritmo predeterminado. Si la semilla aleatoria de Torch se establece en un valor fijo, debería ser posible asegurar que la salida de la misma entrada sea fija cada vez que se ejecute la red.
El segundo es cudnnGetConvolutionForwardAlgorithm_v7 vs cudnnFindConvolutionForwardAlgorithmEx (el enlace es a la API oficial proporcionada por NVIDIA). Ya sabemos que PyTorch llama al primero de forma predeterminada, y establecer benchmark = True llamará al segundo. Según la explicación oficial, la función Get utilizará algunos métodos heurísticos establecidos artificialmente para seleccionar el algoritmo más adecuado considerado por el programa. La función Buscar es una búsqueda exhaustiva, es decir, atraviesa todas las convoluciones opcionales para comparar. En otras palabras, de hecho, PyTorch preseleccionará el algoritmo de convolución para cada capa por defecto. La velocidad es más rápida, pero el resultado seleccionado no es tan bueno. El mecanismo de selección específico no es muy claro (no se encontró información relevante) .
Prueba experimental
Ver a algunos estudiantes aquí todavía puede resultar sospechoso. Ahora que hemos hablado de la teoría y analizado el código fuente, todavía estamos a un paso, así que comparemos los resultados. No se preocupe, venga de inmediato.
Primero, echemos un vistazo al resultado predeterminado. El tamaño de entrada es (32, 3, 224, 224). El valor específico se genera aleatoriamente y no afecta el resultado. El modelo de prueba utilizado es ResNet-101, la GPU es GTX 1060, y el resultado a continuación es El tiempo representa el tiempo de propagación hacia adelante y hacia atrás de un mini-lote:
Model: ResNet-101
Device: cuda
Use CUDNN Benchmark: False
Number of runs: 100
Batch size: 32
1.2628223896026611
0.690316915512085
0.739039421081543
0.7383503913879395
0.7389490604400635
...
0.7488319873809814
0.7504653930664062
0.7499253749847412
Veamos la hora en las mismas condiciones después de configurar torch.backends.cudnn.benchmark = True:
Model: ResNet-101
Device: cuda
Use CUDNN Benchmark: True
Number of runs: 100
Batch size: 32
3.634748935699463
0.5547430515289307
0.6138713359832764
0.6144607067108154
0.6148972511291504
...
0.6201446056365967
0.6197457313537598
0.6202619075775146
De lo anterior, podemos ver que la velocidad efectivamente ha aumentado, y es bastante obvio, alrededor de un 15% más rápido (por supuesto, diferentes redes, diferentes plataformas de hardware tienen diferentes efectos). Lo anterior es el tiempo de cada mini-lote, puedes imaginar el tiempo que puede ahorrar cada época. Y aparte del tiempo extra que se dedica la primera vez, casi no hay efectos secundarios.
Ahora probémoslo de nuevo, ¿qué sucede si cambia la forma de la imagen de entrada? El siguiente ejemplo genera aleatoriamente 5 conjuntos de entradas de diferentes tamaños y prueba el rendimiento de cudnn.benchmark = True:
Model: ResNet-101
Device: cuda
Use CUDNN Benchmark: True
Number of runs: 100
Batch size: 32
Number of scenes: 5
iteration 0 torch.Size([32, 3, 154, 154]) time: 3.30
iteration 0 torch.Size([32, 3, 80, 80]) time: 1.92
iteration 0 torch.Size([32, 3, 116, 116]) time: 2.12
iteration 0 torch.Size([32, 3, 118, 118]) time: 0.57
iteration 0 torch.Size([32, 3, 184, 184]) time: 2.67
iteration 1 torch.Size([32, 3, 154, 154]) time: 0.30
iteration 1 torch.Size([32, 3, 80, 80]) time: 0.16
iteration 1 torch.Size([32, 3, 116, 116]) time: 0.20
iteration 1 torch.Size([32, 3, 118, 118]) time: 0.21
iteration 1 torch.Size([32, 3, 184, 184]) time: 0.43
iteration 2 torch.Size([32, 3, 154, 154]) time: 0.35
iteration 2 torch.Size([32, 3, 80, 80]) time: 0.15
iteration 2 torch.Size([32, 3, 116, 116]) time: 0.20
iteration 2 torch.Size([32, 3, 118, 118]) time: 0.21
iteration 2 torch.Size([32, 3, 184, 184]) time: 0.43
...
Y el resultado por defecto es:
Model: ResNet-101
Device: cuda
Use CUDNN Benchmark: False
Number of runs: 100
Batch size: 32
Number of scenes: 5
iteration 0 torch.Size([32, 3, 191, 191]) time: 2.50
iteration 0 torch.Size([32, 3, 121, 121]) time: 0.39
iteration 0 torch.Size([32, 3, 208, 208]) time: 0.54
iteration 0 torch.Size([32, 3, 205, 205]) time: 0.57
iteration 0 torch.Size([32, 3, 185, 185]) time: 0.48
iteration 1 torch.Size([32, 3, 191, 191]) time: 0.47
iteration 1 torch.Size([32, 3, 121, 121]) time: 0.26
iteration 1 torch.Size([32, 3, 208, 208]) time: 0.54
iteration 1 torch.Size([32, 3, 205, 205]) time: 0.57
iteration 1 torch.Size([32, 3, 185, 185]) time: 0.48
iteration 2 torch.Size([32, 3, 191, 191]) time: 0.47
iteration 2 torch.Size([32, 3, 121, 121]) time: 0.26
iteration 2 torch.Size([32, 3, 208, 208]) time: 0.54
iteration 2 torch.Size([32, 3, 205, 205]) time: 0.57
iteration 2 torch.Size([32, 3, 185, 185]) time: 0.48
...
De acuerdo con los resultados anteriores, podemos saber que cuando cudnn.benchmark = True, si sigue cambiando la forma de entrada, la eficiencia operativa será realmente baja, porque para cada nueva escena convolucional PyTorch encontrará automáticamente el algoritmo de convolución más adecuado. Sin embargo, de acuerdo con el código fuente y los resultados anteriores, también podemos ver que para una determinada escena de convolución optimizada, se guardará en el caché y el resultado se leerá directamente del caché la próxima vez que lo encuentre. . Por el contrario, en el caso predeterminado de PyTorch (es decir, cudnn.benchmark = False), el cambio en el tamaño de entrada no afecta la eficiencia.
Algunos estudiantes respondieron que después de usar la prueba de código en el apéndice, encontraron que el efecto de mejora de la velocidad no es muy obvio. La razón puede ser que la GPU utilizada es mejor y la velocidad de entrenamiento en sí es muy rápida. Puede que no sea obvio después de configurar cudnn.benchmark = True. Por el contrario, debido a que la GPU que uso es relativamente promedio, la diferencia de velocidad es obvia. Para los estudiantes con mejores GPU, puede aumentar apropiadamente el tamaño del lote en el programa y probar nuevamente, o simplemente probar el efecto con el modelo de detección o segmentación a mano. Normalmente, habrá un aumento de velocidad.
para resumir
Este artículo habla principalmente sobre:
一种可以在底层优化卷积层，进而在不改变所使用的卷积神经网络本身结构的情况下提高效率的方法；torch.backends.cudnn.benchmark 如何使用；浏览了一遍 PyTorch 的 cuDNN 代码中有关卷积算法选择的部分；测试了 torch.backends.cudnn.benchmark 的实际使用效果。
Con respecto a cuándo es apropiado establecer torch.backends.cudnn.benchmark = True, una oración es: si la estructura de la red convolucional no cambia dinámicamente y la entrada de la red (tamaño de lote, tamaño de imagen, canal de entrada) es fija, entonces Rest ¡seguro!
apéndice
Parte del código de prueba que utilicé es el siguiente:
import time
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
import numpy as np
parser = argparse.ArgumentParser(description='Test for cudnn.benchmark')
parser.add_argument('--run_num', type=int, default=100, help='number of runs')
parser.add_argument('--batch_size', type=int, default=32, help='batch size')
parser.add_argument('--use_gpu', dest='use_gpu', action='store_true', default=False, help='use gpu')
parser.add_argument('--use_benchmark', dest='use_benchmark', action='store_true', default=False, help='use benchmark')
parser.add_argument('--exp_name', type=str, default='cudnn_test', help='output file name')
args = parser.parse_args()
if args.use_gpu and torch.cuda.is_available():
device = torch.device('cuda')
print('Using GPU: ', torch.cuda.get_device_name(0))
if args.use_benchmark:
torch.backends.cudnn.benchmark = True
print('Using cudnn.benchmark.')
else:
device = torch.device('cpu')
print('Warning! Using CPU.')
images = torch.randn(args.batch_size, 3, 224, 224)
labels = torch.empty(args.batch_size, dtype=torch.long).random_(1000)
model = models.resnet101()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()
model = model.to(device)
images = images.to(device)
labels = labels.to(device)
time_list = []
model.train()
for itr in range(args.run_num):
start = time.time()
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
end = time.time()
print('iteration %d time: %.2f' % (itr, end-start))
time_list.append(end-start)
with open(args.exp_name, 'w') as f:
f.write('Device: ' + device.type + '\n')
f.write('Use CUDNN Benchmark: ' + str(torch.backends.cudnn.benchmark) + '\n')
f.write('Number of runs: ' + str(args.run_num) + '\n')
f.write('Batch size: ' + str(args.batch_size) + '\n')
f.write('Average time: %.2f s\n\n' % (np.mean(time_list)))
for each in time_list:
f.write(str(each))
f.write('\n')