Dataloader dictionary에서 batch를 밖으로 빼내기

Problem

기존에 pytorch로 구현되어 있는 maskrcnn training 코드를 pytorch lightning으로 바꾸다가 발생한 문제.

모델의 loss를 계산하기위해 입력하는 데이터의 형식이 약간 다른 문제가 발생. batch가 dictionary 안에 들어가 있어서 이를 dictionary 밖으로 빼내야 한다.

# 현재 상태
{
    'boxes': tensor([[[ 25., 209., 343., 610.]],
    [[181., 138., 550., 608.]]], device='cuda:1')
}
        
        
# 되어야하는 상태
# batch가 dictionary 밖으로 나와야함
[
    {'boxes': tensor([[ 25., 209., 343., 610.]], device='cuda:1')},
    {'boxes': tensor([[181., 138., 550., 608.]], device='cuda:1')}
]

 

Solution

기존 pytorch에서는

y = [{k: v for k, v in t.items()} for t in targets]

위와 같은 코드로 해결을 했었는데, 잘 동작하지 않아 새로 맞게 짰다.

y = []
for idx in range(len(x)):
    dict = {}
    for key in self.keys:
        dict[key] = targets[key][idx]
    y.append(dict)
    
'''
[
    {
	'boxes': tensor([[ 25., 209., 343., 610.]], device='cuda:1'), 
	'labels': tensor([1], device='cuda:1'), 
	'masks': tensor([[[0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         ...,
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0]]], device='cuda:1', dtype=torch.uint8), 
 	'image_id': tensor([1], device='cuda:1'), 
	'area': tensor(100, device='cuda:1'), 
	'iscrowd': tensor([0], device='cuda:1')
    }, 
    {
	'boxes': tensor([[181., 138., 550., 608.]], device='cuda:1'), 
	'labels': tensor([1], device='cuda:1'), 
	'masks': tensor([[[0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         ...,
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0]]], device='cuda:1', dtype=torch.uint8), 
 	'image_id': tensor([4], device='cuda:1'), 
	'area': tensor(100, device='cuda:1'), 
	'iscrowd': tensor([0], device='cuda:1')}
 ]
'''