【Pytorch】pixelごとにことなる重みをつける方法

要約

• PytorchでLossのbackwardを行うにはLossがスカラーである必要がある。
• torch.nn.BCELoss()は通常平均値を与えるためパラメータでreduction='none'と指定することでピクセルごとのLossを得る
• その後にpixelごとのlayer weightを掛け、torch.mean()としてやることでbackwardが成立する

背景

前からpytorchを使ってsegmentationで遊んでいたけれど、ピクセルごとにweightを掛けたlossを実装とすると、"grad can be implicitly created only for scalar outputs"とエラーが出てしまっていてうまく行かなかった なんとなく放置していたけれど、いよいよやらざるを得なくなってきていたので、本格的に調べて解決した。

方法

要約とソースコード参照

参考URL

• Pytorch公式
• 質問

テストコード

デフォルトのBCEloss(reduction='mean')

x = torch.randn(2, 3, 4, 5, requires_grad=True)
y = torch.randn(2, 3, 4, 5)
criterion = torch.nn.BCEWithLogitsLoss(reduction='mean')
loss = criterion(x, y)
loss.backward()

pixel-wiseで重みをつけようとした場合

x = torch.randn(2, 3, 4, 5, requires_grad=True)
y = torch.randn(2, 3, 4, 5)
w = torch.randn(2, 3, 4, 5) # weight
w[0,0,0,0] = 0
a = torch.ones(x.shape)
criterion = torch.nn.BCEWithLogitsLoss(reduction='none')
loss = criterion(x, y)
loss = loss * w
loss = torch.mean(loss)
loss.backward()