量子アニーリング・イジングマシンによるブラックボックス最適化¶

import os
import torch
import numpy as np


def seed_everything(seed=0):
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

from amplify.client import FixstarsClient
    
client = FixstarsClient()
client.parameters.timeout = 1000  # タイムアウト1秒
# client.token = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"  # ローカル環境等で使用する場合は、Amplify AEのアクセストークンを入力してください。

import torch.nn as nn
    
    
class TorchFM(nn.Module):
    def __init__(self, d: int, k: int):
        super().__init__()
        self.V = nn.Parameter(torch.randn(d, k), requires_grad=True)
        self.lin = nn.Linear(d, 1)  # 右辺第1項及び2項は全結合ネットワーク

    def forward(self, x):
        out_1 = torch.matmul(x, self.V).pow(2).sum(1, keepdim=True)
        out_2 = torch.matmul(x.pow(2), self.V.pow(2)).sum(1, keepdim=True)
        out_inter = 0.5 * (out_1 - out_2)
        out_lin = self.lin(x)
        out = out_inter + out_lin
        return out

from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split

import copy


def train(
    X,
    y,
    model_class=None,
    model_params=None,
    batch_size=1024,
    epochs=3000,
    criterion=None,
    optimizer_class=None,
    opt_params=None,
    lr_sche_class=None,
    lr_sche_params=None,
):
    X_tensor, y_tensor = (
        torch.from_numpy(X).float(),
        torch.from_numpy(y).float(),
    )
    indices = np.array(range(X.shape[0]))
    indices_train, indices_valid = train_test_split(
        indices, test_size=0.2, random_state=42
    )

    train_set = TensorDataset(X_tensor[indices_train], y_tensor[indices_train])
    valid_set = TensorDataset(X_tensor[indices_valid], y_tensor[indices_valid])
    loaders = {
        "train": DataLoader(train_set, batch_size=batch_size, shuffle=True),
        "valid": DataLoader(valid_set, batch_size=batch_size, shuffle=False),
    }

    model = model_class(**model_params)
    best_model_wts = copy.deepcopy(model.state_dict())
    optimizer = optimizer_class(model.parameters(), **opt_params)
    if lr_sche_class is not None:
        scheduler = lr_sche_class(optimizer, **lr_sche_params)
    best_score = 1e18
    for epoch in range(epochs):
        losses = {"train": 0.0, "valid": 0.0}

        for phase in ["train", "valid"]:
            if phase == "train":
                model.train()
            else:
                model.eval()

            for batch_x, batch_y in loaders[phase]:
                optimizer.zero_grad()
                out = model(batch_x).T[0]
                loss = criterion(out, batch_y)
                losses[phase] += loss.item() * batch_x.size(0)

                with torch.set_grad_enabled(phase == "train"):
                    if phase == "train":
                        loss.backward()
                        optimizer.step()

            losses[phase] /= len(loaders[phase].dataset)

        with torch.no_grad():
            model.eval()
            if best_score > losses["valid"]:
                best_model_wts = copy.deepcopy(model.state_dict())
                best_score = losses["valid"]
        if lr_sche_class is not None:
            scheduler.step()

    with torch.no_grad():
        model.load_state_dict(best_model_wts)
        model.eval()
    return model

def gen_training_data(D: int, N0: int, true_func):
    assert N0 < 2**D
    # N0個の入力値を乱数を用いて取得
    X = np.random.randint(0, 2, size=(N0, D))
    # 取得した入力値のうち重複しているものを除外し、除外した分の入力値を乱数を用いて追加
    X = np.unique(X, axis=0)
    while X.shape[0] != N0:
        X = np.vstack((X, np.random.randint(0, 2, size=(N0 - X.shape[0], D))))
        X = np.unique(X, axis=0)
    y = np.zeros(N0)
    # N0個の入力値に対応する出力値を目的関数を評価して取得
    for i in range(N0):
        if i % 10 == 0:
            print(f"Generating {i}-th training data set.")
        y[i] = true_func(X[i])
    return X, y

from amplify import (
    Solver,
    BinarySymbolGenerator,
    sum_poly,
    BinaryMatrix,
    BinaryQuadraticModel,
)
import matplotlib.pyplot as plt
import sys


class FMQA:
    def __init__(self, D: int, N: int, N0: int, k: int, true_func, solver) -> None:
        assert N0 < N
        self.D = D
        self.N = N
        self.N0 = N0
        self.k = k
        self.true_func = true_func
        self.solver = solver
        self.y = None

    # 教師データに基づいて N-N0 回のFMQAを教師データを追加しながら繰り返し実施するメンバー関数
    def cycle(self, X, y, log=False) -> np.ndarray:
        print(f"Starting FMQA cycles...")
        pred_x = X[0]
        pred_y = 1e18
        for i in range(self.N - self.N0):
            print(f"FMQA Cycle #{i} ", end="")
            try:
                x_hat = self.step(X, y)
            except RuntimeError:
                sys.exit(f"Unknown error, i = {i}")
            # x_hat として既に全く同じ入力が教師データ内に存在する場合、その周辺の値を x_hat とする。
            is_identical = True
            while is_identical:
                is_identical = False
                for j in range(i + self.N0):
                    if np.all(x_hat == X[j, :]):
                        change_id = np.random.randint(0, self.D, 1)
                        x_hat[change_id.item()] = 1 - x_hat[change_id.item()]
                        if log:
                            print(f"{i=}, Identical x is found, {x_hat=}")
                        is_identical = True
                        break
            # hat{x} で目的関数 f() を評価
            y_hat = self.true_func(x_hat)
            # 最適点近傍における入出力ペア [x_hat, y_hat] を教師データに追加
            X = np.vstack((X, x_hat))
            y = np.append(y, y_hat)
            # 目的関数の評価値が最小値を更新したら、その入出力ペアを [pred_x, pred_y] へコピー
            if pred_y > y_hat:
                pred_y = y_hat
                pred_x = x_hat
                print(f"variable updated, {pred_y=}")
            else:
                print("")
            # 全ての入力を全探索済みの場合は、for文を抜ける
            if len(y) >= 2**self.D:
                print(f"Fully searched at {i=}. Terminating FMQA cycles.")
                break
        self.y = y
        return pred_x

    # 1回のFMQAを実施するメンバー関数
    def step(self, X, y) -> np.ndarray:
        # FM を機械学習
        model = train(
            X,
            y,
            model_class=TorchFM,
            model_params={"d": self.D, "k": self.k},
            batch_size=8,
            epochs=2000,
            criterion=nn.MSELoss(),
            optimizer_class=torch.optim.AdamW,
            opt_params={"lr": 1},
        )
        # 学習済みモデルから、FM パラメータの抽出
        v, w, w0 = list(model.parameters())
        v = v.detach().numpy()
        w = w.detach().numpy()[0]
        w0 = w0.detach().numpy()[0]
        # ここから量子アニーリング・イジングマシンによる求解を実施
        gen = BinarySymbolGenerator()  # BinaryPoly の変数ジェネレータを宣言
        q = gen.array(self.D)  # BinaryPoly から決定変数の作成
        cost = self.__FM_as_QUBO(q, w0, w, v)  # FM パラメータから QUBO として FM を定義
        result = self.solver.solve(cost)  # 目的関数を Amplify のソルバーに受け渡し
        if len(result.solutions) == 0:
            raise RuntimeError("No solution was found.")
        values = result.solutions[0].values
        q_values = q.decode(values)
        return q_values

    # FM パラメータから QUBO として FM を定義する関数。前定義の TorchFM クラスと同様に、g(x) の関数形通りに数式を記述。
    def __FM_as_QUBO(self, x, w0, w, v):
        lin = w0 + (x.T @ w)
        D = w.shape[0]
        out_1 = sum_poly(self.k, lambda i: sum_poly(D, lambda j: x[j] * v[j, i]) ** 2)
        # 次式において、x[j] はバイナリ変数なので、x[j] = x[j]^2 であることに注意。
        out_2 = sum_poly(
            self.k, lambda i: sum_poly(D, lambda j: x[j] * v[j, i] * v[j, i])
        )
        return lin + (out_1 - out_2) / 2

    """上記の __FM_as_QUBO で用いられている sum_poly は、計算速度やメモリの観点から非効率。一般的に決定変数の相互作用項が非ゼロである FM の場合、BinaryMatrix を使う次の書き方が効率的。ここで、BinaryMatrixでの2次項は、上三角行列で表される非対角項に対応するため、FM式の2次の項に対する x(1/2) は不要。また、上の __FM_as_QUBO（sum_poly を使う実装）と関数のシグネチャを合わせるために、x を引数に取っているが、BinaryMatrix を使う本実装では本来は不要。
    def __FM_as_QUBO(self, x, w0, w, v):
        out_1_matrix = v @ v.T
        out_2_matrix = np.diag((v * v).sum(axis=1))
        matrix = BinaryMatrix(out_1_matrix - out_2_matrix + np.diag(w))
        # 定数項 w0 を忘れずに BinaryQuadraticModel の2つ目の引数に入れる。
        model = BinaryQuadraticModel(matrix, w0)
        return model
    """

    # 初期教師データ及び各 FMQA サイクル内で実施した i 回の目的関数評価値の履歴をプロットする関数
    def plot_history(self):
        assert self.y is not None
        fig = plt.figure(figsize=(6, 4))
        plt.plot(
            [i for i in range(self.N0)],
            self.y[: self.N0],
            marker="o",
            linestyle="-",
            color="b",
        )  # 初期教師データ生成時の目的関数評価値（ランダム過程）
        plt.plot(
            [i for i in range(self.N0, self.N)],
            self.y[self.N0 :],
            marker="o",
            linestyle="-",
            color="r",
        )  # FMQA サイクル時の目的関数評価値（FMQA サイクル過程）
        plt.xlabel("i-th evaluation of f(x)", fontsize=18)
        plt.ylabel("f(x)", fontsize=18)
        plt.tick_params(labelsize=18)
        return fig

# d次元対称行列であって成分の平均が0のものを出力
def make_Q(d) -> np.ndarray:
    Q_true = np.random.rand(d, d)
    Q_true = (Q_true + Q_true.T) / 2
    Q_true = Q_true - np.mean(Q_true)
    return Q_true


# 乱数シード値を初期化
seed_everything(0)

# 入力値の次元（問題サイズ）
D = 100
# 真の関数で使われる行列 Q
Q = make_Q(D)
# 目的関数（xQx）の定義。


def true_func(x):
    # 本来は、未知関数（シミュレーションや実験）の結果値やその後処理結果値を cost とする。
    cost = x @ Q @ x
    return cost


N = 70  # 関数を評価できる回数
N0 = 60  # 初期教師データのサンプル数
k = 10  # FMにおけるベクトルの次元（ハイパーパラメータ）

# client：先に作成した Amplify クライアント
solver = Solver(client)
# 初期教師データの生成
X, y = gen_training_data(D, N0, true_func)

# FMQA のインスタンス化
fmqa_solver = FMQA(D, N, N0, k, true_func, solver)
# FMQA サイクルの実行
pred_x = fmqa_solver.cycle(X, y)
# 最適化結果の出力
print("pred x:", pred_x)
print("pred value:", true_func(pred_x))

fig = fmqa_solver.plot_history()

# BinaryPoly の変数ジェネレータを宣言
gen = BinarySymbolGenerator()
# サイズ D を有する決定変数の1次元配列を作成
q = gen.array(D)
# xQx を QUBO の目的関数として定式化
cost = sum_poly(D, lambda i: sum_poly(D, lambda j: Q[i, j] * q[i] * q[j]))
# 目的関数を Amplify に渡し、求解を行う。
result = solver.solve(cost)
if len(result.solutions) == 0:
    raise RuntimeError("No solution was found.")
# 推定された最適解を抽出し表示する。
values = result.solutions[0].values
true_x = q.decode(values)
print("true x:", true_x)
print("true value:", true_func(true_x))