def range(start, end, initial_step, cache_type=DEFAULT_CACHE_TYPE, context=None): context = context or tst.get_simple_traceback(1) def execute_range(func): return explicit.range(func, start, end, initial_step, toma_cache_type=cache_type, toma_context=context) return execute_range
def batch(initial_batchsize, cache_type=DEFAULT_CACHE_TYPE, context=None): context = context or tst.get_simple_traceback(1) def execute_batch(func): return explicit.batch(func, initial_batchsize, toma_cache_type=cache_type, toma_context=context) return execute_batch
def get_batchsize(self, initial_batchsize: int): stacktrace = tst.get_simple_traceback(2) if self.TRACK_RAM: cpu_available_memory = int( toma.cpu_memory.get_available_cpu_memory() // self.MEMORY_GRANULARITY) else: cpu_available_memory = -1 gpu_available_memory = int(tcm.get_cuda_assumed_available_memory() // self.MEMORY_GRANULARITY) batchsize = self.get_batchsize_from_cache(stacktrace, cpu_available_memory, gpu_available_memory) batchsize.set_initial_batchsize(initial_batchsize) return batchsize
def chunked( tensor: torch.Tensor, initial_step: Optional[int] = None, dimension: Optional[int] = None, cache_type: Type = DEFAULT_CACHE_TYPE, context=None, ): context = context or tst.get_simple_traceback(1) def execute_chunked(func): return explicit.chunked( func, tensor, initial_step, toma_dimension=dimension, toma_cache_type=cache_type, toma_context=context, ) return execute_chunked
def get_stacktrace(): return stacktrace.get_simple_traceback()
def __init__(self): stacktrace = tst.get_simple_traceback(2) BatchsizeCache.all_instances[stacktrace] = self