pytorch class weight 주는법

pytorch에서 모델을 학습할 때 특정 클래스에 가중치를 주어 더 잘 학습하게 하는 방법이 있습니다.

weights = torch.FloatTensor([1/20, 1/20, 1/20, 7/20, 1/20, 7/20, 1/20, 1/20, 1/100]).cuda()
criterion = torch.nn.CrossEntropyLoss(weight=weights)

위의 상황은 클래스가 9개일 때 인덱스 3과 5번은 더 잘 탐지하고 싶고, 마지막 인덱스는 불필요한 상황입니다.

이후 아래와 같이 설정해주면 됩니다.

logits = model(inputs)
labels = labels.view(bs*seq_length)

loss = criterion(logits, labels)