[Pytorch] ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 256, 1, 1]) (feat. drop_last)

 

상황

1. model.train()인 상태로 dataset을 학습시키고

2. model 내에는 batch norm layer가 존재할 경우

3. dataset을 batch size만큼씩 나눌 때 가끔씩 나머지가 1인 경우가 존재하는데 이때 발생할 수 있는 에러이다.

(ex. dataset의 수가 65개인데 batch size를 64인 경우). 

 

에러

ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 256, 1, 1])

 

원인

batch norm layer의 계산식에 batch chanel의 크기가 1보다 크다고 가정한 수식이 존재하는 듯하다. 따라서, batch size가 1일 경우에는 따로 처리해주어야 한다.

 

해결방법

해결방법은 dataloader 부분에 drop_last=True의 파라미터를 추가시켜주면 해결된다. 이는 batch의 마지막은 버리도록 해주는 코드이다. 예를 들면, 데이터의 양이 40개일 때 batch_size를 16으로 해주면 16 / 16 / 8 이렇게 3개의 epoch를 돌아야 하는데 8은 버리는 방식이다. 이로 인하여 마지막 batch에 담겨 버려지는 데이터의 정보는 학습하지 못한다는 단점이 존재하지만 데이터의 크기가 커지면 영향이 적기에 default 값으로 넣어주는 것이 좋다. 

import torch

train_loader = torch.utils.data.DataLoader(
    dataset         = train_data,
    batch_size      = BATCH_SIZE,
    drop_last       = True
)

 

잘 해결되었다.

 

 

 

-히비스서커스-