在Pythorch中编写适应可用（CUDA）内存的算法
toma的Python项目详细描述
火炬记忆自适应算法（TOMA）
帮助程序的集合，以便于编写适应可用（CUDA）内存的代码。 特别是，它会重试由于OOM（内存不足）而失败的代码，并自动降低批处理大小。在
为了避免重复故障转移，我们实现了一个简单的缓存，它在给定调用和可用空闲内存的情况下，存储上次成功的批大小。在
安装
要使用pip安装，请使用：
pip install toma
要运行测试，请使用：
^{pr2}$示例
fromtomaimporttoma@toma.batch(initial_batchsize=512)defrun_inference(batchsize,model,dataset):# ...run_inference(batchsize,model,dataset)
这将尝试执行batchsize=512的train_模型。如果抛出内存错误，它将减小批处理大小，直到成功为止。在
注意：
此批大小可以与用于累积渐变的批大小不同，方法是每隔一段时间调用optimizer.step()。在
为了更容易在一个范围内循环，还有toma.range和{
@toma.chunked(initial_step=512)defcompute_result(out:torch.Tensor,start:int,end:int):# ...result=torch.empty((8192,...))compute_result(result)
这将分块result，并将这些块逐个传递给compute_result。
同样，如果由于OOM而失败，步骤将被半等分。
与toma.batch相比，这允许在块上循环时减少步长。
这样可以节省计算。在
@toma.range(initial_step=32)defreduce_data(start:int,end:int,out:torch.Tensor,dataA:torch.Tensor,dataB:torch.Tensor):# ...reduce_data(0,1024,result,dataA,dataB)
{{{cd8> 为了更容易地执行一个块，而不必将其提取到函数中然后调用它，我们还提供了 目前有3种可用的缓存类型。
可以通过设置 例如： 或者 这将存储给定调用跟踪的成功批大小和该点的可用内存。
对于大多数机器学习代码，这足以记住正确的batchsize，而不必查看实际参数和理解更多的语义。在 隐含的假设是，在几次迭代之后，GPU和CPU内存使用率将达到一个稳定状态。在 为了限制进程的CPU内存，toma提供了： 这也有助于避免意外的交换抖动。在 这将独立于调用发生的位置重用上一个成功的batchsize。在 总是从建议的批处理大小开始，并在必要时进行故障转移。在 这涉及到开销。Toma只能与其他耗时/内存消耗的操作一起使用。在 感谢@y0ast的反馈和讨论。在 标签： 欢迎加入QQ群-->： 979659372
toma.executetoma.execute.batch、toma.execute.range和{deffunction():# ... other code@toma.execute.chunked(batched_data,initial_step=128):defcompute(chunk,start,end):# ...
缓存
toma.DEFAULT_CACHE_TYPE或将cache_type传递给调用来更改它们。在@toma.batch(initial_batchsize=512,cache_type=toma.GlobalBatchsizeCache)
toma.explicit.batch(...,toma_cache_type=toma.GlobalBatchsizeCache)
StacktraceMemoryBatchsizeCache:Stacktrace&Available Memory（默认值）importtoma.cpu_memorytoma.cpu_memory.set_cpu_memory_limit(8)
GlobalBatchsizeCache：全局每函数NoBatchsizeCache：无缓存基准/间接费用
---------------------------------------------------------------------------------- benchmark: 5 tests ----------------------------------------------------------------------------------
Name (time in ms) Min Max Mean StdDev Median IQR Outliers OPS Rounds Iterations
----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_native 2.1455 (1.0) 3.7733 (1.0) 2.3037 (1.0) 0.1103 (1.0) 2.2935 (1.0) 0.1302 (1.0) 81;5 434.0822 (1.0) 448 1
test_simple 17.4657 (8.14) 27.0049 (7.16) 21.0453 (9.14) 2.6233 (23.79) 20.4881 (8.93) 3.4384 (26.42) 13;0 47.5165 (0.11) 39 1
test_toma_no_cache 31.4380 (14.65) 40.8567 (10.83) 33.2749 (14.44) 2.2530 (20.43) 32.2698 (14.07) 2.8210 (21.67) 4;1 30.0527 (0.07) 25 1
test_explicit 33.0759 (15.42) 52.1866 (13.83) 39.6956 (17.23) 6.9620 (63.14) 38.4929 (16.78) 11.2344 (86.31) 4;0 25.1917 (0.06) 20 1
test_toma 36.9633 (17.23) 57.0220 (15.11) 43.5201 (18.89) 6.7318 (61.05) 41.6034 (18.14) 7.2173 (55.45) 2;2 22.9779 (0.05) 13 1
----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
谢谢
推荐PyPI第三方库
