Hello everybody,
I’m trying to reduce the inference of my system by running two halves of the model in parallel. I’m using DETR object detection.
Basically, I have two threads running in parallel, one runs the backbone and another one runs the detection part of the network. However in order for this to work, I need to store the feature map resulting from the backbone on the previous iteration. This is highly increasing my memory usage by 4 times. I don’t understand why this is 4 times higher, In my understanding it should only be two times higher, since at some point I will have both the tensors from the previous iteration and the tensors from the current iteration.
I’m printing the memory allocated (2178970624) and the memory reserved (4309647360), and there is a two times difference between them.
Here is some of the code (If some part of the code is missing it is because i had to remove a lot of lines to post it here)
def run_backbone2(model,data_queue, backbone_results_queue):
while True:
samples, ind, targets = data_queue.get(block=True)
src,mask, pos = model.run_backbone(samples)
del samples
# torch.cuda.empty_cache()
# backbone_results_queue.put((src.cpu(),mask.cpu(), pos.cpu(), ind, targets))
backbone_results_queue.put((src, mask, pos, ind, targets))
def run_detection2(model, backbone_results_queue,outputs_queue):
while True:
src,mask, pos, ind, targets = backbone_results_queue.get(block=True)
# outputs = model.run_detection(src.cuda(),mask.cuda(), pos.cuda())
outputs = model.run_detection(src, mask, pos)
del src
del mask
del pos
# torch.cuda.empty_cache()
outputs_queue.put((outputs, ind, targets))
@torch.no_grad()
def evaluate_parallel(model, criterion, postprocessors, dataset, device, output_dir,
virtual_epoch_len=100):
model.cuda()
model.eval()
criterion.eval()
data_queue = Queue(1)
backbone_results_queue = Queue(1)
output_queue = Queue(1)
backbone_process = threading.Thread(target= run_backbone2, args=(model,data_queue,
backbone_results_queue))
backbone_process.start()
detection_process = threading.Thread(target=run_detection2,
args=(model, backbone_results_queue, output_queue))
detection_process.start()
for _ind in trange(virtual_epoch_len, desc='COCO evaluation: '):
samples, _targets = dataset.__getitem__(_ind)
samples = samples.unsqueeze(0)
samples = samples.to(device)
if _ind > 0:
outputs, ind, targets = output_queue.get(block = True)
del outputs
_targets = [_targets]
_targets = [{k: v.to(device) for k, v in t.items()} for t in _targets]
data_queue.put((samples, _ind, _targets))
print(torch.cuda.memory_allocated())
print(torch.cuda.memory_reserved())
