写了才会，所以直接上代码运行，代码参考自零基础AI入门指南 - Blogs - 廖雪峰的官方网站，下述操作都是在Windows下，需预先安装Python。
1.PyTorch训练框架安装
pip install torch torchvision torchaudio
安装界面如下：
PyTorch训练框架安装
2.编写模型
import torch.nn as nn
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.fc1 = nn.Linear(in_features=64 * 5 * 5, out_features=128)
self.fc2 = nn.Linear(in_features=128, out_features=10)
def forward(self, x):
x = nn.functional.relu(self.conv1(x))
x = nn.functional.max_pool2d(x, kernel_size=2)
x = nn.functional.relu(self.conv2(x))
x = nn.functional.max_pool2d(x, kernel_size=2)
x = x.view(-1, 64 * 5 * 5)
x = nn.functional.relu(self.fc1(x))
x = self.fc2(x)
return x
将上述代码拷贝到一个txt文件，并重命名为model.py。
3.训练
from time import time
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from model import NeuralNetwork
def train(dataloader, device, model, loss_fn, optimizer):
model.train()
running_loss = 0.0
for batch, (inputs, labels) in enumerate(dataloader):
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'loss: {running_loss/len(dataloader):>0.3f}')
def test(dataloader, device, model):
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in dataloader:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'accuracy: {100.0*correct/total:>0.2f} %')
def main():
print('loading training data...')
train_data = datasets.MNIST(
root='./data', train=True, download=True, transform=ToTensor())
print('loading test data...')
test_data = datasets.MNIST(
root='./data', train=False, download=True, transform=ToTensor())
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'using {device}')
model = NeuralNetwork().to(device)
print(model)
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
epochs = 5
for t in range(epochs):
start_time = time()
print(f'epoch {t+1} / {epochs}\n--------------------')
train(train_dataloader, device, model, loss_fn, optimizer)
test(test_dataloader, device, model)
end_time = time()
print(f'time: {end_time-start_time:>0.2f} seconds')
print('done!')
path = 'mnist.pth'
torch.save(model.state_dict(), path)
print(f'model saved: {path}')
if __name__ == '__main__':
main()
model.py同级目录另外新建一个txt文件，将上述代码拷贝到次文件中，并重命名为trainer.py，注意trainer.py有调用到model.py，运行trainer.py：
训练
训练会生成一个文件mnist.pth。
4.使用模型
4.1 准备图片
用ps手绘10张png图片，图片分辨率最好设置为28*28，太大会导致预测效果很差，10张图片按"test-序号.png"命名，序号从0开始，
4.2 测试代码
在trainer.py新建一个文件test.py，并将下述代码拷贝到此文件中：
import torch
from torchvision import transforms
from PIL import Image, ImageOps
from model import NeuralNetwork
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'using {device}')
model = NeuralNetwork().to(device)
path = './mnist.pth'
model.load_state_dict(torch.load(path))
print(f'loaded model from {path}')
print(model)
def test(path):
print(f'test {path}...')
image = Image.open(path).convert('RGB').resize((28, 28))
image = ImageOps.invert(image)
trans = transforms.Compose([
transforms.Grayscale(1),
transforms.ToTensor()
])
image_tensor = trans(image).unsqueeze(0).to(device)
model.eval()
with torch.no_grad():
output = model(image_tensor)
probs = torch.nn.functional.softmax(output[0], 0)
predict = torch.argmax(probs).item()
return predict, probs[predict], probs
def main():
for i in range(10):
predict, prob, probs = test(f'./input/test-{i}.png')
print(f'expected {i}, actual {predict}, {prob}, {probs}')
if __name__ == '__main__':
main()
- 注意准备的图片命名序号要从0开始，并将图片保存到input目录中，如果不够10张，可按下图方式修改测试图片数量
调整数量
-
运行成功效果如下：
测试
5.服务部署
新建文件server.py，并将下述代码拷贝到此文件：
import base64
import torch
from io import BytesIO
from PIL import Image
from flask import Flask, request, redirect, jsonify
from torchvision import transforms
from model import NeuralNetwork
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'using {device}')
model = NeuralNetwork().to(device)
path = './mnist.pth'
model.load_state_dict(torch.load(path))
print(f'loaded model from {path}')
print(model)
params = model.state_dict()
print(params)
app = Flask(__name__)
@app.route('/')
def index():
return redirect('/static/index.html')
@app.route('/api', methods=['POST'])
def api():
data = request.get_json()
image_data = base64.b64decode(data['image'])
image = Image.open(BytesIO(image_data))
trans = transforms.Compose([
transforms.Grayscale(1),
transforms.ToTensor()
])
image_tensor = trans(image).unsqueeze(0).to(device)
model.eval()
with torch.no_grad():
output = model(image_tensor)
probs = torch.nn.functional.softmax(output[0], 0)
predict = torch.argmax(probs).item()
prob = probs[predict]
print(f'predict: {predict}, prob: {prob}, probs: {probs}')
return jsonify({
'result': predict,
'probability': prob.item()
})
if __name__ == '__main__':
app.run(port=5000)
- 运行该服务需要安装flask，在控制台输入以下命令安装即可：
pip install Flask
-
在控制台输入python server.py，启动服务：
启动服务
6.静态页面运行
从参考链接对应github上将静态网页文件夹static拷贝过来，然后打开浏览器在地址栏输入：
127.0.0.1:5000/static/index.html
运行效果如下：
实测效果
7.总结
总的来说，构建一个自己的AI训练框架，需要经过构建模型、训练得到训练生成文件，模型加载生成文件，然后喂图片数据，得到输出。
