高周波エンジニアのためのAI・機械学習入門（GPU編5）PythonとKeras3.0を使って3次のLCバンドパスフィルタ（BPF）のSパラメータからTransformerで素子の値（L、C）推定する。LSTMを使うより10倍くらい速くなった。

LSTMが遅いのでTransformerが開発されたような話を聞いたことがある。ということで今回はTransformer。

といってもTransformerは言語モデルで使う話ばかりで回帰で使う例があまり出てこない…のでChatGPTと相談しながらコードを書く。

コードはこちら。

import os

os.environ["KERAS_BACKEND"] = "jax"

import keras

import jax

import numpy as np

import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split

from sklearn.preprocessing import StandardScaler

plt.rcParams['font.family'] = 'Noto Sans CJK JP'

# JAXの高速化

jax.config.update("jax_enable_x64", False)

jax.config.update("jax_disable_jit", False)

# mixed precision (JAX + Keras3対応)

keras.mixed_precision.set_global_policy("mixed_float16")

# ==========================

# データ読み込み

# ==========================

data_label = np.load("data_label.npz")

data = data_label["data"]     # (N,200,5)

label = data_label["label"]   # (N,6)

x_train, x_test, y_train, y_test = train_test_split(

    data, label, test_size=0.3, random_state=0

)

# ==========================

# ラベルのスケーリング

# ==========================

scaler_y = StandardScaler()

y_train_f = scaler_y.fit_transform(y_train)

y_test_f  = scaler_y.transform(y_test)

# ==========================

# Positional Encoding

# ==========================

def positional_encoding(seq_len, d_model):

    pos = np.arange(seq_len)[:, np.newaxis]

    i = np.arange(d_model)[np.newaxis, :]

    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))

    angle_rads = pos * angle_rates

    pe = np.zeros((seq_len, d_model))

    pe[:, 0::2] = np.sin(angle_rads[:, 0::2])

    pe[:, 1::2] = np.cos(angle_rads[:, 1::2])

    return keras.ops.convert_to_tensor(pe, dtype="float32")

# ==========================

# Transformer Encoder Block

# ==========================

def transformer_encoder(x, head_size, num_heads, ff_dim, dropout=0.1):

    # Self Attention

    attn_out = keras.layers.MultiHeadAttention(

        num_heads=num_heads,

        key_dim=head_size

    )(x, x)

    #attn_out = keras.layers.Dropout(dropout)(attn_out)

    x = keras.layers.LayerNormalization(epsilon=1e-6)(x + attn_out)

    # Feed Forwardネットワーク

    ffn = keras.layers.Dense(ff_dim, activation="relu")(x)

    ffn = keras.layers.Dense(x.shape[-1])(ffn)

    #ffn = keras.layers.Dropout(dropout)(ffn)

    x = keras.layers.LayerNormalization(epsilon=1e-6)(x + ffn)

    return x

# ==========================

# モデル定義

# ==========================

seq_len = 200

feat_dim = 5

embed_dim = 64

inputs = keras.Input(shape=(seq_len, feat_dim))

# 入力を高次元に射影

x = keras.layers.Dense(embed_dim)(inputs)

# Positional Encodingを加算

pe = positional_encoding(seq_len, embed_dim)

x = x + pe

# Transformer ブロック

x = transformer_encoder(x, head_size=64, num_heads=4, ff_dim=128)

#x = transformer_encoder(x, head_size=64, num_heads=4, ff_dim=128)

# 時系列次元を集約

x = keras.layers.GlobalAveragePooling1D()(x)

# 全結合

x = keras.layers.Dense(128, activation="relu")(x)

outputs = keras.layers.Dense(6)(x)

model = keras.Model(inputs, outputs)

model.compile(

    optimizer=keras.optimizers.Adam(learning_rate=1e-3),

    loss="mean_squared_error"

)

model.summary()

# ==========================

# 学習

# ==========================

batch_size = 32

epochs = 300

keras.utils.set_random_seed(1)

history = model.fit(

    x_train,

    y_train_f,

    batch_size=batch_size,

    epochs=epochs,

    validation_split=0.15,

)

# ==========================

# 予測

# ==========================

y_pred_f = model.predict(x_test)

y_pred = scaler_y.inverse_transform(y_pred_f)

# ==========================

# R2評価

# ==========================

metric = keras.metrics.R2Score()

metric.update_state(y_test, y_pred)

print("\nR2 Score:", metric.result())

# ==========================

# %誤差

# ==========================

eps = 1e-12

valid_mask = np.abs(y_test) > eps

pct_error = np.zeros_like(y_test)

pct_error[valid_mask] = np.abs(

    (y_test[valid_mask] - y_pred[valid_mask]) / y_test[valid_mask] * 100

)

print("\n% Error per element (mean):", pct_error.mean(axis=0))

# ==========================

# 相関プロット

# ==========================

row, column = 2, 3

legend = ["L1", "C1", "L2", "C2", "L3", "C3"]

fig, ax = plt.subplots(2, 3, figsize=(15, 9))

for i in range(row):

    for j in range(column):

        count = column * i + j

        maxvalue = y_pred[:, count].max()

        ax[i,j].scatter(y_pred[:, count], y_test[:,count], c="r", s=5)

        ax[i,j].plot([0, maxvalue], [0, maxvalue], "--", c="black")

        ax[i,j].set_xlabel("推定した値")

        ax[i,j].set_ylabel("実際の値")

        ax[i,j].set_xlim(0, maxvalue)

        ax[i,j].set_ylim(0, maxvalue)

        ax[i,j].grid()

        ax[i,j].legend([legend[count] + f" 平均誤差{pct_error.mean(axis=0)[count]:.2f}%"])

fig.tight_layout()

plt.show()

モデルは以下のようになる。Transformer層は最初2つ重ねていたが、余計悪くなった。1つでよさそう。

結果はこちら。CNN並みの精度でLSTMより10倍以上速い。

とはいえこの事例だとやっぱり速度的にCNNが一番よさそう。