Pytorchで独自の損失関数を設計したときにgrad_fnが無いと言われた

背景

既存モデルのLossをマルチラベルに対応できるように改良していたときに以下で詰まっていた．

lement 0 of tensors does not require grad and does not have a grad_fn

やりたいこととしては多クラス分類→多ラベル分類． そのときに元々あったLossであるnn.CrossEntropyLoss をシグモイド出力からのnn.BCEWithLogitsLoss()に変えたかった．

原因

おそらくlossの値を逆伝播するときにgrad_fnを参考に行っていると考えられる． 独自の平均だったり合計などをlistとかに変えて計算しまうと，これらの情報が消えているのでおかしくなるのでは？

grad_fnはどんな計算があってそのテンソルができたかを残しておくもので，例えば平均によって計算されたら以下になる．

tensor(0.7343, device='cuda:0', grad_fn=<MeanBackward0>)

解決

torch.stackでlistを変換してmeanを適応した．

rtn = torch.stack(losses).mean()

ちなみにlossesにはBCEWithLogitsLoss()で計算されたlossの値が入っている配列．

[tensor(0.5282, device='cuda:0',
       grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), tensor(0.9063, device='cuda:0',
       grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), tensor(0.7409, device='cuda:0',
       grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), tensor(0.6982, device='cuda:0',
       ...

また，結果をtensor型で新たに定義して引数requires_grad=Trueを設定（以下コード）することでもできるが，ここまで保持してた他情報（例えばGPUのdeviceとか）がなくなるので面倒．

torch.tensor(losses,requires_grad=True)

---

例

まず擬似的なマルチラベルと最終層で計算されたベクトルを作る．

import random 
label_mtx = torch.tensor([[random.choice([0,1]) for i in range(6)] for j in range(20)], dtype=torch.float32,device=device,requires_grad=True)
score_mtx = torch.tensor([[random.random() for i in range(6)] for j in range(20)], dtype=torch.float32,device=device,requires_grad=True)

そしてLossを計算する関数を定義. LossはBCEWithLogitsLoss()， （ちなみにBCEWithLogitsLossは内部でシグモイドを適応してくれるみたいなのでモデルの出力をそのまま入れて良いらしい）

def mtx_BCEWithLogitsLoss(score_mtx,label_mtx):
    loss_fn = nn.BCEWithLogitsLoss()
    losses = []
    for score_i,label_i in zip(score_mtx,label_mtx):
        loss_i = loss_fn(score_i,label_i)
        losses.append(loss_i)

    rtn = torch.stack(losses).mean()
    return rtn

実行すると以下．

mtx_BCEWithLogitsLoss(score_mtx,label_mtx)
>>>tensor(0.7343, device='cuda:0', grad_fn=<MeanBackward0>)

ちゃんとtensorに情報を持っています．

参考にしたもの

Kaggleにあった画像マルチラベルしているやつ www.kaggle.com

もし違う場合など指摘があればコメントして欲しいです．