RedisAI: 深度学习部署的新选择

RedisAI简介

RedisAI是Redis实验室和[tensor]werk合作开发的一个Redis模块,旨在将机器学习(ML)和深度学习(DL)带入Redis生态系统,使AI模型能够尽可能靠近数据存储的位置执行。这个项目的目标是为深度学习模型的生产部署提供一个简单而强大的解决方案。

RedisAI允许用户在Redis中存储张量(多维数组)、深度学习模型和执行脚本。它支持多种主流深度学习框架,包括TensorFlow、PyTorch和ONNX,使开发者能够无缝集成他们现有的AI/ML工作流。

RedisAI架构

RedisAI的主要特性

张量存储: RedisAI引入了新的数据结构来存储多维数组(张量),这些张量可以直接用作模型的输入或输出。
模型服务: 支持加载和执行来自TensorFlow、PyTorch和ONNX的预训练模型。
脚本执行: 通过TorchScript提供了在CPU或GPU上执行预处理和后处理操作的能力。
DAG执行: 允许定义和执行包含多个操作的有向无环图(DAG),实现复杂的推理管道。
GPU加速: 支持GPU执行,充分利用硬件加速能力。
扩展性: 利用Redis集群和Redis Sentinel,RedisAI可以轻松扩展到多节点部署。

为什么选择RedisAI?

与传统的深度学习部署方案相比,RedisAI提供了以下优势:

简化部署: 对于已经使用Redis的团队来说,RedisAI可以轻松集成到现有架构中,无需引入新的复杂组件。
性能优化: 通过将模型执行放在数据附近,RedisAI减少了数据移动,提高了推理性能。
灵活性: 支持多种深度学习框架,允许开发者使用他们熟悉的工具。
可扩展性: 利用Redis的分布式特性,RedisAI可以轻松扩展以处理大规模工作负载。
简化运维: 对于已经熟悉Redis运维的团队,RedisAI不需要学习新的复杂系统。

使用RedisAI

要开始使用RedisAI,您可以通过Docker快速设置一个开发环境:

docker pull redisai/redisai
docker run -p 6379:6379 -it --rm redisai/redisai

或者,您可以从源代码构建RedisAI:

git clone https://github.com/RedisAI/RedisAI.git
cd RedisAI
bash get_deps.sh
mkdir build
cd build
cmake -DDEPS_PATH=../deps/install ..
make

构建完成后,您可以使用以下命令启动Redis服务器并加载RedisAI模块:

redis-server --loadmodule /path/to/redisai.so

RedisAI实战:图像分类示例

让我们通过一个简单的图像分类示例来展示RedisAI的强大功能。在这个例子中,我们将使用预训练的MobileNet模型来识别图像中的对象。

首先,我们需要设置预处理脚本:

def pre_process(image):
    image = image.float()
    image /= 255
    return image

将此脚本保存到RedisAI:

import redisai as rai
con = rai.Client(host='localhost', port=6379)
with open('script.txt') as f:
    script = f.read()
con.scriptset('preprocess', rai.Device.cpu, script)

加载预训练的MobileNet模型:

with open('mobilenet.pb', 'rb') as f:
    model = f.read()
con.modelset('mobilenet', rai.Backend.tf, rai.Device.cpu, model,
             inputs=['input'], outputs=['MobilenetV2/Predictions/Reshape_1'])

执行推理管道:

import numpy as np
from PIL import Image

# 加载和预处理图像
image = Image.open('panda.jpg')
image_tensor = rai.Tensor.from_numpy(np.array(image).astype(np.float32))
con.tensorset('image', image_tensor)

# 执行预处理脚本
con.scriptrun('preprocess', 'pre_process', inputs=['image'], outputs=['processed_image'])

# 运行模型
con.modelrun('mobilenet', inputs=['processed_image'], outputs=['output'])

# 获取结果
result = con.tensorget('output')

解析结果:

import json

with open('imagenet_classes.json') as f:
    classes = json.load(f)

top_k = result.argmax(axis=1)[0][:5]
for i in top_k:
    print(f"{classes[str(i)]}: {result[0][i]*100:.2f}%")

这个例子展示了如何使用RedisAI来执行一个完整的深度学习推理管道,包括数据预处理、模型执行和结果后处理。

RedisAI的未来发展

RedisAI团队正在积极开发新功能,以进一步增强其功能和性能:

ONNX Runtime支持: 这将允许运行从几乎所有主流深度学习框架导出的模型,包括传统机器学习算法。
自动批处理: 提高大规模推理的效率。
动态后端加载: 允许在运行时动态加载和卸载AI后端。
与其他Redis模块集成: 例如与RediSearch和RedisTimeSeries的集成,为更复杂的AI应用场景提供支持。
RedisAI Enterprise: 为企业级部署提供更多的管理和监控功能。

结论

RedisAI为深度学习模型的生产部署提供了一个强大而灵活的解决方案。通过将AI能力直接集成到Redis生态系统中,它简化了部署过程,提高了性能,并为开发者提供了极大的灵活性。无论您是在构建计算机视觉应用、自然语言处理系统,还是任何需要高性能AI推理的应用,RedisAI都可能成为您技术栈中的重要组成部分。

随着AI技术的不断发展和应用场景的日益复杂,RedisAI将继续进化,为开发者提供更多强大的工具和功能。我们期待看到更多创新的AI应用借助RedisAI的力量得以实现和部署。

如果您正在寻找一种简单而高效的方式来部署您的深度学习模型,RedisAI无疑值得一试。它不仅能够满足您当前的需求,还能为未来的扩展和创新提供坚实的基础。立即开始使用RedisAI,探索AI部署的新可能性吧!