libffcv/ffcv

Reuse memory?

erow opened this issue · 0 comments

erow commented

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