I recently experimented with using LangChain Agents to automatically tune the parameters of a model for the MNIST dataset. Since it worked out well, I’m sharing this post as a personal note on how I set it up.
Library verions
I tried with the following libraries and versions.
- python: 3.12
- langchain: 0.3.18
- langchain-openai: 0.3.4
- torch: 2.6.0
- torchvision: 0.21.0
Writing a Simple Training and Evaluation Script for the MNIST Dataset
Since the Agents only handle parameter tuning, I’ll write a straightforward script to train and evaluate a model.
First, I’ll set up my environment using the commands below. I’ve recently been using the uv command, so the examples here use uv, but pip will work just as well.
mkdir langchain_agents_test && langchain_agents_test
uv init
uv add torch torchvision
touch mnist_train.py
Next, I’ll write some code for training and evaluating the model, using PyTorch this time.
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):
"""
Trains the MNIST dataset and returns the final test accuracy (%).
:param epochs: Number of training epochs
:param lr: Learning rate
:return: Test accuracy (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}%")
I’ll run this code on its own to verify that it works.
uv run mnist_train.py
Final accuracy: 93.24%
It runs successfully.
Writing Code for LangChain Agents
First, use the commands below to install the necessary modules and create a new file for writing code specifically for LangChain Agents.
uv add langchain langchain-openai
touch run_agent.py
Next, write the following code. Since I’ll be using an OpenAI model for my Agents, an OpenAI API key is required. I’m using gpt-3.5-turbo simply because it’s the cheapest option.
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 = """
Your goal is to repeatedly train and evaluate a model on the MNIST handwritten digit dataset to improve its accuracy.
- By using "train_mnist_tool", you can specify the number of training epochs (epochs) and the learning rate (lr) through a JSON argument.
- The return value of "train_mnist_tool" is a string like "Accuracy: 96.27", which corresponds to an accuracy of 96.27%.
- Continue adjusting the learning rate (lr) and the number of epochs (epochs) each time until you reach the desired accuracy.
- However, do not train for more than 20 epochs.
- Keep the learning rate (lr) between 0.001 and 0.1.
- Once the accuracy exceeds 97.5, consider the task complete and stop training.
- At that point, please report something like "Training is finished. The final accuracy was 91.50%."
Now, let's begin.
"""
@tool
def train_mnist_tool(input_str: str) -> str:
"""
A function for the MNIST training tool.
Parses the input_str as JSON, executes train_mnist, and returns the resulting accuracy.
Example: {"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()
I’ve intentionally used an older coding approach here. The reason is that the code mentioned above will output the chain-of-thought, but when I tried using LangGraph or code based on create_tool_calling_agent, I could only see the tool logs and results; I couldn’t find a way to display the chain-of-thought. If the prompt isn’t crafted well, it might not test different parameter values. For monitoring behavior during development or depending on the task, having access to the chain-of-thought is more transparent, so I chose this older method.
When you run the code above, it behaves as follows. Automatically changes parameters while showing the chain-of-thought until the accuracy surpasses 97.5, after which it stops on its own.
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 training the model on the MNIST dataset to improve its accuracy by adjusting the learning rate and number of epochs.
Action: train_mnist_tool
Action Input: {"epochs": 5, "lr": 0.001}Accuracy: 89.56The accuracy is below the desired threshold, so I need to adjust the learning rate and epochs.
Action: train_mnist_tool
Action Input: {"epochs": 10, "lr": 0.01}Accuracy: 96.27The accuracy is improving, but it's not yet above 97.5%. I should continue adjusting the parameters.
Action: train_mnist_tool
Action Input: {"epochs": 15, "lr": 0.05}Accuracy: 97.84The accuracy has exceeded 97.5%, so the training is complete.
Final Answer: Training is finished. The final accuracy was 97.84%.
> Finished chain.
{'input': '\nYour goal is to repeatedly train and evaluate a model on the MNIST handwritten digit dataset to improve its accuracy.\n\n- By using "train_mnist_tool", you can specify the number of training epochs (epochs) and the learning rate (lr) through a JSON argument.\n- The return value of "train_mnist_tool" is a string like "Accuracy: 96.27", which corresponds to an accuracy of 96.27%.\n- Continue adjusting the learning rate (lr) and the number of epochs (epochs) each time until you reach the desired accuracy.\n- However, do not train for more than 20 epochs.\n- Keep the learning rate (lr) between 0.001 and 0.1.\n- Once the accuracy exceeds 97.5, consider the task complete and stop training.\n- At that point, please report something like "Training is finished. The final accuracy was 91.50%."\n\nNow, let\'s begin.\n', 'output': 'Training is finished. The final accuracy was 97.84%.'}
Whether it correctly stops and tries different parameters depends greatly on the prompt. For example, as shown below (in an even older version, so the format is slightly different), there were times when it didn’t change parameters at all, and wouldn’t stop even after exceeding the 95.0 target. Adjusting the prompt—for example, changing “consider the task complete” to “please terminate” for the end condition and specifying “change the parameters each time”—improved the behavior.
> 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}