toma:帮助您在PyTorch中编写适用于可用（CUDA）内存的算法

火炬记忆自适应算法（TOMA） 一组帮助程序，可以更轻松地编写适应于可用（CUDA）内存的代码。具体来说，它将重试由于OOM（内存不足）条件而失败的代码，并自动降低批处理大小。 为了避免重复故障转移，实现了一个简单的缓存，该缓存会在给定调用和可用空闲内存的情况下记住上一次成功的批处理大小。 安装 要使用pip进行安装，请使用： pip install toma 要运行测试，请使用： python setup.py test 例子 from toma import toma @ toma . batch ( initial_batchsize = 512 ) def run_inference ( batchsize , model , dataset ): # ... run_inference ( batchsize , model , dataset ) 这将尝试以batc