高周波エンジニアのためのAI・機械学習入門（GPU編24）TDRでは出せない複素特性インピーダンスの周波数特性を予測する。今回はkerasとバックエンドにJaxを使ったTransformerで。

DNN,CNN,RNN(LSTM)とやってきて今回はTransformer。

コードはこちら。

import numpy as np

import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split

from sklearn.preprocessing import StandardScaler

import os

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

import keras

from jax import config

config.update("jax_enable_x64", False)

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

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

data = data_label["data"]

label = data_label["label"]

x_train, x_test, y_train, y_test = train_test_split(data, label, test_size=0.3, random_state=0)

y_train_list = [y_train[:,:,1], y_train[:,:,2]]

y_test_list = [y_test[:,:,1], y_test[:,:,2]]

scaler_y = [StandardScaler() for _ in range(2)]

y_train_list_f = [None for _ in range(2)]

for i in range(2):

    scaler_y[i] = StandardScaler()

    y_train_list_f[i] = scaler_y[i].fit_transform(y_train_list[i])

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

# 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(200)(x) for i in range(2)]

model = keras.Model(inputs=inputs, outputs=outputs)

model.compile(loss = 'mean_squared_error' ,optimizer=keras.optimizers.Adam())

batch_size = 128

epochs = 1000

keras.utils.set_random_seed(1)

history = model.fit(

    x_train,

    y_train_list_f,

    batch_size=batch_size,

    epochs=epochs,

    validation_split=0.15,

)

y_pred_list_f = model.predict(x_test)

y_pred_list = [None for _ in range(2)]

for i in range(2):

    y_pred_list[i] = scaler_y[i].inverse_transform(y_pred_list_f[i])

metric = keras.metrics.R2Score()

for idx in range(2):

    print(f"R2 score for label {idx}:")

    metric.update_state(y_test_list[idx], y_pred_list[idx])

    result = metric.result()

    print(result)

R2スコアはこちら。まずまず。

R2 score for label 0:
0.99998623
R2 score for label 1:
0.9999977

代表例。ちょっとずれてるか。

最大誤差をヒストグラムにしたもの。

そこそこいいが、今回の例ではCNNが一番精度、速度ともによさそう。