PyTorch でloss関数にtorch.sqrtを用いる方法について

Question

前提・実現したいこと

pytorchでdeep learningの画像分類モデルを作成しています。

モデルのloss関数の一部にtorch.sqrt()をしようしたところ、backward時にnanが発生する問題に突き当たりました。
torch.sqrt()に入力されるベクトルの要素の大きさがとても小さいことが原因のようです。

torch.sqrt()のinputが小さいとbackward時に1/(2*torch.sqrt())がinfになるようです...

何か対処法がわかる方がいらっしゃいましたら、お教えいただければ幸いです。

エラーメッセージ
Traceback (most recent call last):
  File "main_label_grad.py", line 504, in <module>
    model_g = main()
  File "main_label_grad.py", line 459, in main
    tr_acc, tr_acc5, tr_los, grad_train, last_v4 = train(train_loader, net_c, net_t, optimizer_c, optimizer_t, epoch, args, log_G, args.noise_dim, grad_train_old=None, v4=None)
  File "main_label_grad.py", line 320, in train
    loss_trans.backward()
  File "C:\Users\GUESTUSER.conda\envs\tf37\lib\site-packages\torch\tensor.py", line 118, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "C:\Users\GUESTUSER.conda\envs\tf37\lib\site-packages\torch\autograd\__init__.py", line 93, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function 'SqrtBackward' returned nan values in its 0th output.

上記のloss_transがモデルの目的関数で、以下の関数の1つ目の返り値に該当します。
以下の関数(new_norm)の return torch.sqrt(v4_ema)のv4_emaが小さくてnanになってしまっております。

def new_norm(v, epoch, iter, last_v4=None):
    v2 = torch.pow(v,2)
    v4 = torch.pow(v,4)
    v4_ema = ema(v4, epoch, iter, last_v4)
    epsilon = torch.ones(v4_ema.size(0)) * 1e-10
    epsilon = epsilon.cuda()
    return (v2/(torch.sqrt(v4_ema)+epsilon)).sum()/v4_ema.size(0), v4_ema

また上記の関数new_normは以下の式を求めようとして作成したものです。

Comments

[quickquip]

マルチポスト https://ja.stackoverflow.com/questions/59692/
"また、解決した際には必ずteratail及びすべての投稿に解決した旨と、どのように解決したかを記載してください。" https://teratail.com/help#posted-otherservice

情報の追加などを求められた時両方に対応しないといけなくて、結局は面倒なだけで、いいことは何もないですよ。

[fujifuji_]

解決した際や情報の追加を求められた際はは、両方のサイトにお答えさせていただく、所存です。

[fujifuji_]

stackoverflowにおいて、https://ja.stackoverflow.com/questions/59692/pytorchでdeep-learningの画像分類モデルを作成しています

「infになってしまう時のinputの値は完全に0でしょうか？分数の時にinfとならないように分母に十分小さな値(例えば1e-5)を足しておいても同じエラーになりますか？」このような質問をいただきました

[fujifuji_]

「現状のプログラムでは、ルート内部に入る値が1e-13とかになっており、1e-5を足してしまうと、プログラムが目的通りに作用しないと考えております。」と回答いたしました

[fujifuji_]

また、1e-5を足すとプログラムは一応動きました。