機械学習・AI
32 view

LangChain Agents を使用して自動的に機械学習モデルのパラメータをチューニングさせる

LangChain の Agents の勉強用に、MNIST データセットに対するモデルのパラメータを自動的にチューニングさせてみたらそれっぽいことができたので、方法をメモしておきます。

ライブラリのバージョン

以下のバージョンで試しました。

python: 3.12
langchain: 0.3.18
langchain-openai: 0.3.4
torch: 2.6.0
torchvision: 0.21.0

普通に MNIST データセットに対して学習と評価するコードを書く

Agents に実施させるのはパラメータのチューニングのみなので、モデルの学習・評価をするコードは普通に書きます。

まずは適当に以下のコマンドで環境を整備します。最近 uv を使っているので、uv でやっていますが、pip でも問題ありません。

mkdir langchain_agents_test && langchain_agents_test
uv init
uv add torch torchvision
touch mnist_train.py

次にモデルの学習・評価をするコードを以下のように書きます。今回は PyTorch を使用しました。

import random
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms


seed = 0
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True


def train_mnist(epochs=1, lr=0.01):
    """
    MNISTを学習して、最終的なテスト精度(%)を返す関数
    :param epochs: 学習エポック数
    :param lr: 学習率
    :return: テスト精度 (float)
    """
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1000, shuffle=False)

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.fc1 = nn.Linear(28*28, 128)
            self.fc2 = nn.Linear(128, 10)

        def forward(self, x):
            x = x.view(-1, 28*28)
            x = F.relu(self.fc1(x))
            x = self.fc2(x)
            return x

    model = Net()
    optimizer = optim.SGD(model.parameters(), lr=lr)

    # Training
    for epoch in range(epochs):
        model.train()
        for data, target in train_loader:
            optimizer.zero_grad()
            output = model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer.step()

    # Testing
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
            total += len(target)

    accuracy = 100.0 * correct / total
    return accuracy


if __name__ == "__main__":
    acc = train_mnist(epochs=3, lr=0.01)
    print(f"Final accuracy: {acc:.2f}%")

一旦、このコード単体で動かして動くか確認します。

uv run mnist_train.py 
Final accuracy: 93.24%

正常に動作しました。

LangChain Agents のコードを書く

以下のコマンドで必要なモジュールをインストールと LangChain Agents 用のコードを書くために新しいファイルを作成します。

uv add langchain langchain-openai
touch run_agent.py

次に以下のコードを書きます。今回は Agents に使う LLM モデルに OpenAI のものを使うので、OpenAI の API キーを取得してください。一番安いという理由で gpt-3.5-turbo を使っています。

import json
from langchain import hub
from langchain.agents import AgentExecutor, create_react_agent
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool

from mnist_train import train_mnist


OPENAI_API_KEY = "..."
INSTRUCTION = """
あなたはMNISTの手書き数字データセットに対して、
モデルの学習・評価を繰り返し実行しながら精度を向上させることが目的です。

- "train_mnist_tool"を使うと、JSON引数を渡すことで学習エポック数epochsや学習率lrを指定できます
- "train_mnist_tool"の戻り値は文字列で "Accuracy: 96.27" のように返ってきますがこれは96.27%の精度を意味します
- 望みの精度に到達するまで、学習率lrやエポック数epochsを毎回変更して何度でも学習を試してください
- ただし20エポック以上は学習しないでください
- lrは0.001~0.1の範囲で変更してください
- 最終的にAccuracyが97.5を超えたらタスク完了とみなして終了してください
- その際は「学習を終了します。最終的な精度は91.50%でした」のように報告してください

では、はじめてください。
"""


@tool
def train_mnist_tool(input_str: str) -> str:
    """
    MNIST学習Tool用の関数。
    input_strをJSONとして解析し、train_mnistを実行した結果(Accuracy)を返す。
    例: {"epochs": 5, "lr": 0.001}
    """
    try:
        params = json.loads(input_str)
        epochs = params.get("epochs", 1)
        lr = params.get("lr", 0.01)

        accuracy = train_mnist(epochs=epochs, lr=lr)
        return f"Accuracy: {accuracy:.2f}"
    except Exception as e:
        return f"Error: {str(e)}"


def main() -> None:
    prompt = hub.pull("hwchase17/react")

    model = ChatOpenAI(
        model_name="gpt-3.5-turbo",
        temperature=0,
        openai_api_key=OPENAI_API_KEY
    )

    tools = [train_mnist_tool]

    agent = create_react_agent(
        model,
        tools=tools,
        prompt=prompt
    )

    agent_executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        max_iterations=10
    )

    response = agent_executor.invoke({"input": INSTRUCTION})
    print(response)


if __name__ == "__main__":
    main()

コードはあえて古い書き方をしています。その理由は、上記のコードだと思考内容も出力してくれるのですが、LangGraph や create_tool_calling_agent を使うコードだと、Tool を使ったログと結果は出力されるものの、思考のログを出力する方法が自分は見つけられなかったからです。プロンプトをうまく書かないと、パラメータを変えて試さなかったりしたので、動作確認している間やタスクによっては思考のログが出力される方が状況がわかりやすいため、古い書き方にしました。

上記のコードを実行すると以下のようになりました。自動的に Accuracy が 97.5 を超えるまで思考内容を出力しつつパラメータを変えながら試してくれて、超えたら自動的に終了してくれています。

uv run run_agent.py
/Users/zuqqhi2/Desktop/work/projects/Sandbox/langchain_agents/.venv/lib/python3.12/site-packages/langsmith/client.py:253: LangSmithMissingAPIKeyWarning: API key must be provided when using hosted LangSmith API
  warnings.warn(


> Entering new AgentExecutor chain...
I need to start by using the train_mnist_tool function to train the model with different epochs and learning rates to improve accuracy.
Action: train_mnist_tool
Action Input: {"epochs": 5, "lr": 0.001}Accuracy: 89.56The accuracy is below the desired threshold, so I need to continue training with different parameters.
Action: train_mnist_tool
Action Input: {"epochs": 10, "lr": 0.01}Accuracy: 96.27The accuracy has improved significantly, but it's not yet above 97.5%. I should continue training with different parameters.
Action: train_mnist_tool
Action Input: {"epochs": 15, "lr": 0.05}Accuracy: 97.84The accuracy has exceeded 97.5%, so I can stop training now.
Final Answer: 学習を終了します。最終的な精度は97.84%でした (I will stop training now. The final accuracy was 97.84%)

> Finished chain.
{'input': '\nあなたはMNISTの手書き数字データセットに対して、\nモデルの学習・評価を繰り返し実行しながら精度を向上させることが目的です。\n\n- "train_mnist_tool"を使うと、JSON引数を渡すことで学習エポック数epochsや学習率lrを指定できます\n- "train_mnist_tool"の戻り値は文字列で "Accuracy: 96.27" のように返ってきますがこれは96.27%の精度を意味します\n- 望みの精度に到達するまで、学習率lrやエポック数epochsを毎回変更して何度でも学習を試してください\n- ただし20エポック以上は学習しないでください\n- lrは0.001~0.1の範囲で変更してください\n- 最終的にAccuracyが97.5を超えたらタスク完了とみなして終了してください\n- その際は「学習を終了します。最終的な精度は91.50%でした」のように報告してください\n\nでは、はじめてください。\n', 'output': '学習を終了します。最終的な精度は97.84%でした (I will stop training now. The final accuracy was 97.84%)'}

終了判定やパラメータをちゃんと変えて試してくれるかは、プロンプトに結構依存する感じでした。以下 (さらに古い書き方で試した時のものなので形式は少し違います) のような感じでパラメータを一切変えず、そして目標値の 95.0 を超えているのに終了しなかったりする時もありました。プロンプトにて、パラメータについて「毎回変更して」や終了条件のところで「タスク完了とみなして」とするだけでなく「終了してください」という表現に変えたりするなど、細かい修正を加えたら改善しました。

> Entering new AgentExecutor chain...
I need to train the MNIST model with different epochs and learning rates until the accuracy exceeds 95.0%.
Action: train_mnist_tool
Action Input: {"epochs": 10, "lr": 0.01}
Observation: Accuracy: 95.94
Thought:The accuracy has exceeded 95.0%, but I should continue training to improve it further.
Action: train_mnist_tool
Action Input: {"epochs": 10, "lr": 0.01}
Observation: Accuracy: 96.47
Thought:The accuracy is still improving, I should continue training.
Action: train_mnist_tool
Action Input: {"epochs": 10, "lr": 0.01}
Observation: Accuracy: 96.17
Thought:The accuracy seems to fluctuate slightly, I will train one more time to see if it improves.
Action: train_mnist_tool
Action Input: {"epochs": 10, "lr": 0.01}
Observation: Accuracy: 96.20
Thought:The accuracy is close to 95.0%, I will train one more time to see if it exceeds the threshold.
Action: train_mnist_tool
Action Input: {"epochs": 10, "lr": 0.01}
Observation: Accuracy: 96.36
Thought:The accuracy is still below 95.0%, I will train one more time to try to reach the threshold.
Action: train_mnist_tool
Action Input: {"epochs": 10, "lr": 0.01}