pytorchの自動微分を関数内で行いたい

Question

pytorchで自動微分を行いたいときに、A1からA5の微分可能なテンソルを用意して、func内で縮約した結果に対してA1からA5で勾配を取得したいのですが、
辞書形式で、下のように関数名を取得すると、requires_gradが無効になってしまいます。

Ai = locals().get(f'A{i+1}', 0)

下記のプログラムのように、縮約計算した後に勾配を得たいときに、動的に配列名を得る方法はありますか？
またA1～A5はグローバル変数として扱いたいです。

import torch
import numpy as np
A1 = torch.arange(4).reshape(2,2).to(torch.float32)
A2 = torch.arange(8).reshape(2,2,2).to(torch.float32)
A3 = torch.arange(8).reshape(2,2,2).to(torch.float32)
A4 = torch.arange(8).reshape(2,2,2).to(torch.float32)
A5 = torch.arange(4).reshape(2,2).to(torch.float32)

A1.requires_grad = True
A2.requires_grad = True
A3.requires_grad = True
A4.requires_grad = True
A5.requires_grad = True

def func(A1,A2,A3,A4,A5):
    I = torch.tensor([0, 1], dtype=torch.float32) 
    A = torch.einsum("ia,i->a",A1,I)
    for i in range(1,4):
        Ai = locals().get(f'A{i+1}', 0)
        A = torch.einsum("a,aib,i->b",A, Ai,I)
    y = torch.einsum("a,ai,i->...",A,A5,I)
    return y

y=func(A1,A2,A3,A4,A5)
y.backward()
print(A2.grad)

Comments

[melian]

for i in range(1,4):
　Ai = locals().get(f'A{i+1}', 0)
　A = torch.einsum("a,aib,i->b",A, Ai,I)

上記のループは locals() を使わず、以下の様にしてもよいのではないでしょうか。(変数の個数は3個なので、そのまま書き下してもよいかと)

for Ai in (A2, A3, A4):
　A = torch.einsum("a,aib,i->b",A, Ai,I)

[ps_aux_grep]

関数func内でlocals()を介して取得しているA2,A3,A4について，requires_gradを確認しましたが全てTrueでした．具体的にどこが問題なのでしょうか

[quickquip]

> 辞書形式で、下のように関数名を取得すると、requires_gradが無効になってしまいます。
はどういう意味ですか?

こういうコードを実行したら、こういう結果になりました
期待する結果はこうでした
というような、解釈が要らない書き方にしてもらえるといいと思います

[TakaiY]

前段階の質問に、A1～A5のような数値付き変数でなくリストや辞書などを使った方がいいという正攻法の回答が付いているのに、それを使わず locals()というようなトリッキーなことをして、わざわざ複雑さを導入するようなことはやめたほうがいいと思います。 おかげで何が問題なのかわかりにくくなっています。

[redcatML]

確認ありがとうございます。
requires_gradはtrueになっていますね。
確認不足で失礼しました。

Answer 1

質問でできていなかった点が、間違いであることが分かりました。