ffcv icon indicating copy to clipboard operation
ffcv copied to clipboard

Reuse memory?

Open erow opened this issue 1 year ago • 0 comments

I found that self.loader.graph.allocate_memory takes a lot of memory. It seems work for reusing the memory.

 >>> ffcv/loader/epoch_iterator.py 

if not hasattr(self.loader,'memory_allocations'):
    print("Memory Allocation 19219.1 MiB for bs=128 and batches_ahead=10")
    self.memory_allocations = self.loader.graph.allocate_memory(
        self.loader.batch_size,
        self.loader.batches_ahead + 2
    )
    self.loader.memory_allocations = self.memory_allocations
else:
    print("Reuse memory")
    self.memory_allocations = self.loader.memory_allocations

erow avatar Feb 12 '24 23:02 erow