PyTorchはPythonの深層学習フレームワークです。この記事では、PINNで2次元キャビティ流れを解くなどの物理シミュレーション応用を念頭に置き、PyTorchの基礎であるテンソル・自動微分・ニューラルネットワークの定義・学習ループを解説します。

テンソルとは何か｜PyTorchの基本データ構造

テンソルの定義

テンソルはPyTorchの基本データ構造で、スカラー・ベクトル・行列・多次元配列を統一的に扱います。NumPyのndarrayに相当しますが、テンソルはGPU上で演算できる点と自動微分に対応している点が異なります。

import torch

# スカラー（0次元テンソル）
a = torch.tensor(3.0)

# ベクトル（1次元テンソル）
b = torch.tensor([1.0, 2.0, 3.0])

# 行列（2次元テンソル）
c = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

print(a.shape)  # torch.Size([])
print(b.shape)  # torch.Size([3])
print(c.shape)  # torch.Size([2, 2])

よく使うメソッド

使いがちな0埋め、1埋め、乱数は以下のとおりです。乱数は一様乱数[0, 1)を生成します。

print(torch.zeros(1, 3)) # tensor([[0., 0., 0.]])
print(torch.ones(1, 2))  # tensor([[1., 1.]])
print(torch.(2, 1))      # tensor([[0.1221],[0.2996]])

形状変換とスライス

reshapeで形状を変換できます。要素数が一致していれば自由に変形できます。

x = torch.zeros(100)　　# torch.Size([100])
x = x.reshape(100, 1)   # 列ベクトルに変換
print(x.shape)           # torch.Size([100, 1])

スライスはNumPyと同じ記法です。

a = torch.tensor([[1.0, 2.0, 3.0],
                   [4.0, 5.0, 6.0]])
print(a[:, 0])    # 1列目: tensor([1., 4.])
print(a[0, :])    # 1行目: tensor([1., 2., 3.])
print(a[:, 0:1])  # 1列目（2次元を保持）: shape [2, 1]

PINNでは座標点の配列を (N, 2) の形状（N点 × x,y成分）で扱います。[:, 0:1] のようにスライスすることで、2次元の形状を保ったまま各成分を取り出せます。

GPUへの転送

.to(device) でテンソルをGPUに転送します。

device = "cuda:0" if torch.cuda.is_available() else "cpu"
x = torch.rand(100, 2).to(device)

自動微分（autograd）の仕組み

requires_grad

PyTorchは計算グラフを動的に構築し、チェーンルール（連鎖律）によって微分を自動計算します。

テンソルに requires_grad=True を設定すると、そのテンソルを含む演算がすべて記録されます。

具体例を見てみましょう。

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x   # y = x^2 + 3x

y.backward()          # dy/dx を計算
print(x.grad)         # tensor(7.)  ← dy/dx|_{x=2} = 2x+3 = 7

backward() を呼ぶと、計算グラフを逆向きにたどって各テンソルの .grad に勾配が蓄積されます。

多次元テンソルへの微分

テンソルの各要素に対して微分できます。この機能を、PINNでは座標点 (x, y) に対してネットワーク出力の偏微分 $\partial u / \partial x$ などを計算するために使います。

coords = torch.tensor([[1.0, 2.0],
                        [3.0, 4.0]], requires_grad=True)

# coords[:, 0] が x 成分、coords[:, 1] が y 成分
f = coords[:, 0] ** 2 + coords[:, 1]   # f = x^2 + y

# f の全要素の和について微分（スカラーにしてからbackward）
f.sum().backward()

print(coords.grad)
# tensor([[2., 1.],
#         [6., 1.]])
# ∂f/∂x = 2x → [2, 6], ∂f/∂y = 1 → [1, 1]

このとき、backward() を呼ぶにはスカラー（単一の値）である必要があります。テンソルがスカラーでない場合、PyTorchはどの成分について微分を計算するかを特定できないためエラーになります。そのため、テンソルの全要素の和を計算してスカラー化することで、backward() を適用可能にしています。考えてみると、損失関数の計算結果ってスカラーですよね。

torch.autograd.grad

backward() はLossの計算に使われたすべてのノードに対して勾配を計算しますが、代わりに torch.autograd.grad を使うと、特定のテンソルに対する微分を明示的に取得できます。PINNのPDE残差計算では、ネットワーク出力を座標で微分する際にこの関数を使います。

x = torch.tensor([[1.0], [2.0], [3.0]], requires_grad=True)
u = x ** 3   # u = x^3

du_dx = torch.autograd.grad(
    outputs=u,
    inputs=x,
    grad_outputs=torch.ones_like(u),
    create_graph=True,   # 高階微分を可能にする
)[0]

print(du_dx)
# tensor([[3.],   ← 3x^2|_{x=1} = 3
#         [12.],  ← 3x^2|_{x=2} = 12
#         [27.]]) ← 3x^2|_{x=3} = 27

create_graph=True を指定すると、微分の計算グラフも保持されるため、さらにその微分（2階微分）を計算できます。PINNでは $\partial^2 u / \partial x^2$ のような2階微分が必要なため、この引数が重要です。

d2u_dx2 = torch.autograd.grad(
    outputs=du_dx,
    inputs=x,
    grad_outputs=torch.ones_like(du_dx),
    create_graph=True,
)[0]

print(d2u_dx2)
# tensor([[6.],   ← 6x|_{x=1} = 6
#         [12.],  ← 6x|_{x=2} = 12
#         [18.]]) ← 6x|_{x=3} = 18

勾配の蓄積とzero_grad

backward() を呼ぶたびに .grad に勾配が加算されます。PINNではPDE残差の計算が必要なので、学習ループでは毎ステップ optimizer.zero_grad() で勾配をリセットする必要があります。これを忘れると勾配が累積し、Lossの計算結果が不適切になります。

ニューラルネットワークの定義方法｜nn.Module

nn.Moduleの基本

PyTorchのニューラルネットワークは nn.Module を継承したクラスとして定義します。__init__ でレイヤーを定義し、forward で順伝播の計算を記述する決まりです。

import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2, 64)    # 入力2次元 → 64次元
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)    # 64次元 → 出力1次元

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

nn.Linear(in, out) は全結合層で、入力 $\mathbf{x} \in \mathbb{R}^{in}$ に対して

$\mathbf{y} = \mathbf{W}\mathbf{x} + \mathbf{b}$

を計算します。$\mathbf{W} \in \mathbb{R}^{out \times in}$ が重み行列、$\mathbf{b} \in \mathbb{R}^{out}$ がバイアスベクトルです。どちらも学習対象のパラメータです。

ここでは紹介しませんが、PyTorchにはほかにも様々な層のクラスが用意されています。

活性化関数

活性化関数は線形レイヤーの間に挟む非線形関数です。活性化関数を挟むおかげでニューラルネットワークは非線形になり、高い表現力を獲得できます。

ニューラルネットワークの学習でよく使われる活性化関数を以下に示します。

関数	式	特徴
ReLU	$\max(0, x)$	計算が速い、勾配消失が起きにくい
tanh	$\tanh(x)$	出力が[-1, 1]に収まる、PINNで使われる
SiLU	$x \cdot \sigma(x)$	滑らか、PINNで良い結果が出ることが多い

$\sigma(x) = 1/(1 + e^{-x})$ はシグモイド関数です。SiLUはSwishとも呼ばれます。

PINNのような2階微分を必要とするタスクでは活性化関数の選定に注意が必要です。例えば、ReLUは2階微分すると必ず0になるため、学習が不安定化したり、損失関数の収束性が悪化することがあります。そのため、tanhやSiLUのような2階微分がゼロにならない活性化関数を選ぶことがポイントです。

# SiLUを使った例
class PinnNet(nn.Module):
    def __init__(self, in_features, out_features, num_layers, layer_size):
        super().__init__()
        self.input_layer = nn.Linear(in_features, layer_size)
        self.hidden_layers = nn.ModuleList(
            [nn.Linear(layer_size, layer_size) for _ in range(num_layers - 1)]
        )
        self.output_layer = nn.Linear(layer_size, out_features)
        self.activation = nn.SiLU()

    def forward(self, x):
        x = self.activation(self.input_layer(x))
        for layer in self.hidden_layers:
            x = self.activation(layer(x))
        x = self.output_layer(x)
        return x

学習ループの書き方｜損失関数・オプティマイザ・パラメータ更新

損失関数

損失関数はモデルの予測値と正解の「ずれ」を数値化する関数です。学習ではこの損失を最小化するようにパラメータを更新します。

回帰問題でよく使われる平均二乗誤差（MSE）は次の式で定義されます。

$\mathcal{L} = \frac{1}{N} \sum_{i=1}^{N} (\hat{y}_i - y_i)^2$

$\hat{y}_i$ はモデルの予測値、$y_i$ は正解値、$N$ はデータ点数です。PyTorchでは torch.mean(...) を使って自分で書くか、nn.MSELoss() を使います。

# 自分で書く場合
loss = torch.mean((pred - target) ** 2)

# nn.MSELossを使う場合
criterion = nn.MSELoss()
loss = criterion(pred, target)

PINNではPDE残差・境界条件残差など複数の損失項を足し合わせた損失関数を使います。PINNによるcavity流れの学習では、MSELossといった既存の損失関数のクラスを使わず torch.mean(residual ** 2) のように定式化しました。

オプティマイザ

オプティマイザはパラメータをどのように更新するかを決めるアルゴリズムです。Adamは学習率を自動で調整する手法で、深層学習では広く使われています。

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

スケジューラ

深層学習では、学習が進むにつれて学習率を下げるのが一般的で、これにより収束性が向上します。学習率のスケジューリングには torch.optim.lr_scheduler.LambdaLR クラスを用います。

scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=lambda step: 0.9999 ** step   # 指数的に減衰
)

学習ループの全体像

num_steps = 10000

for step in range(num_steps):
    # 1. 勾配をリセット
    optimizer.zero_grad()

    # 2. 順伝播（モデルに入力を与えて出力を得る）
    pred = model(x_train)

    # 3. 損失を計算
    loss = torch.mean((PDE_loss) ** 2)

    # 4. 逆伝播（損失をパラメータで微分）
    loss.backward()

    # 5. パラメータを更新
    optimizer.step()
    scheduler.step()

    if step % 1000 == 0:
        print(f"step {step}, loss = {loss.item():.6f}")

モデルの保存と読み込み

PyTorchでは、torch.save でパラメータをファイルに保存し、load_state_dict で読み込めます。

# 保存
torch.save(model.state_dict(), "model.pth")

# 読み込み
model.load_state_dict(torch.load("model.pth", map_location=device))
model.eval()

推論時は model.eval() を呼ぶことで、ドロップアウトやバッチ正規化が評価モードに切り替わります。

さらに深く学ぶための書籍

PyTorchの動作をより深く学びたい方には、PyTorchの公式サイトをおすすめします。