378 lines
12 KiB
Python
378 lines
12 KiB
Python
|
import torch
|
|||
|
|
import time
|
|||
|
|
import os
|
|||
|
|
from PIL import Image
|
|||
|
|
from transformers import AutoImageProcessor, AutoModelForImageClassification
|
|||
|
|
from flask import Flask, request, jsonify
|
|||
|
|
from io import BytesIO
|
|||
|
|
|
|||
|
|
# 引入寒武纪MLU相关模块
|
|||
|
|
try:
|
|||
|
|
import torch_mlu
|
|||
|
|
print(f"成功导入torch_mlu，版本: {getattr(torch_mlu, '__version__', 'unknown')}")
|
|||
|
|
|
|||
|
|
# 简单的MLU可用性测试
|
|||
|
|
def check_mlu_available():
|
|||
|
|
try:
|
|||
|
|
test_tensor = torch.randn(2, 2).mlu()
|
|||
|
|
return True
|
|||
|
|
except:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def get_mlu_device_count():
|
|||
|
|
"""通过尝试多个设备索引来检测可用的MLU设备数量"""
|
|||
|
|
max_devices_to_check = 8
|
|||
|
|
available_devices = 0
|
|||
|
|
|
|||
|
|
for i in range(max_devices_to_check):
|
|||
|
|
try:
|
|||
|
|
test_tensor = torch.randn(2, 2).mlu(i)
|
|||
|
|
available_devices += 1
|
|||
|
|
print(f"MLU设备 {i} 可用")
|
|||
|
|
except:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
return available_devices
|
|||
|
|
|
|||
|
|
def get_device_name(device_index):
|
|||
|
|
"""获取设备名称"""
|
|||
|
|
try:
|
|||
|
|
return f"MLU-Device-{device_index}"
|
|||
|
|
except:
|
|||
|
|
return f"MLU-Device-{device_index} (Unknown)"
|
|||
|
|
|
|||
|
|
# 创建模拟的ct模块
|
|||
|
|
class MLUModel:
|
|||
|
|
@staticmethod
|
|||
|
|
def is_mlu_available():
|
|||
|
|
return check_mlu_available()
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def device_count():
|
|||
|
|
return get_mlu_device_count()
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def get_device_name(device_index):
|
|||
|
|
return get_device_name(device_index)
|
|||
|
|
|
|||
|
|
ct = MLUModel()
|
|||
|
|
MLU_AVAILABLE = check_mlu_available()
|
|||
|
|
print(f"MLU设备可用: {MLU_AVAILABLE}")
|
|||
|
|
print(f"检测到 {ct.device_count()} 个MLU设备")
|
|||
|
|
|
|||
|
|
except ImportError:
|
|||
|
|
torch_mlu = None
|
|||
|
|
ct = None
|
|||
|
|
MLU_AVAILABLE = False
|
|||
|
|
print("警告: 未找到torch_mlu模块，无法使用MLU设备")
|
|||
|
|
except Exception as e:
|
|||
|
|
torch_mlu = None
|
|||
|
|
ct = None
|
|||
|
|
MLU_AVAILABLE = False
|
|||
|
|
print(f"MLU初始化警告: {str(e)}")
|
|||
|
|
|
|||
|
|
# 设置CPU核心数
|
|||
|
|
os.environ["OMP_NUM_THREADS"] = "4"
|
|||
|
|
os.environ["MKL_NUM_THREADS"] = "4"
|
|||
|
|
os.environ["NUMEXPR_NUM_THREADS"] = "4"
|
|||
|
|
os.environ["OPENBLAS_NUM_THREADS"] = "4"
|
|||
|
|
os.environ["VECLIB_MAXIMUM_THREADS"] = "4"
|
|||
|
|
torch.set_num_threads(4)
|
|||
|
|
|
|||
|
|
class MLUImageClassifier:
|
|||
|
|
def __init__(self, model_path: str):
|
|||
|
|
# 检测并使用MLU设备
|
|||
|
|
self.use_mlu = self._check_mlu_availability()
|
|||
|
|
print(f"使用设备: {'MLU' if self.use_mlu else 'CPU'}")
|
|||
|
|
|
|||
|
|
# 加载处理器和模型
|
|||
|
|
self.processor = AutoImageProcessor.from_pretrained(model_path)
|
|||
|
|
self.model = self._load_model(model_path)
|
|||
|
|
self.id2label = self.model.config.id2label
|
|||
|
|
|
|||
|
|
# 验证模型设备
|
|||
|
|
self._verify_model_device()
|
|||
|
|
|
|||
|
|
def _check_mlu_availability(self):
|
|||
|
|
"""检查MLU设备是否可用"""
|
|||
|
|
if torch_mlu is None:
|
|||
|
|
print("MLU不可用: torch_mlu模块未找到")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 测试MLU基本功能
|
|||
|
|
test_tensor = torch.randn(2, 2).mlu()
|
|||
|
|
test_result = test_tensor + test_tensor
|
|||
|
|
print("MLU设备可用性测试通过")
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"MLU设备测试失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def _load_model(self, model_path: str) -> AutoModelForImageClassification:
|
|||
|
|
"""加载模型到合适的设备"""
|
|||
|
|
try:
|
|||
|
|
# 先在CPU加载模型
|
|||
|
|
model = AutoModelForImageClassification.from_pretrained(
|
|||
|
|
model_path,
|
|||
|
|
torch_dtype=torch.float32
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if self.use_mlu:
|
|||
|
|
# 先将模型完全移动到CPU确保稳定
|
|||
|
|
model = model.cpu()
|
|||
|
|
|
|||
|
|
# 使用.mlu()方法将模型移动到MLU设备
|
|||
|
|
model = model.mlu()
|
|||
|
|
print("模型成功加载到MLU设备")
|
|||
|
|
else:
|
|||
|
|
model = model.cpu()
|
|||
|
|
print("模型加载到CPU设备")
|
|||
|
|
|
|||
|
|
return model.eval()
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"模型加载失败: {str(e)}")
|
|||
|
|
# 尝试fallback到CPU模式
|
|||
|
|
try:
|
|||
|
|
model = model.cpu()
|
|||
|
|
print("Fallback到CPU模式")
|
|||
|
|
return model.eval()
|
|||
|
|
except:
|
|||
|
|
raise RuntimeError(f"模型加载完全失败: {str(e)}")
|
|||
|
|
|
|||
|
|
def _verify_model_device(self):
|
|||
|
|
"""验证模型设备"""
|
|||
|
|
try:
|
|||
|
|
param = next(self.model.parameters())
|
|||
|
|
if self.use_mlu:
|
|||
|
|
# 对于MLU设备，通过简单操作验证
|
|||
|
|
test_output = param + 0
|
|||
|
|
print("MLU模型验证成功")
|
|||
|
|
else:
|
|||
|
|
print("CPU模型验证成功")
|
|||
|
|
|
|||
|
|
except StopIteration:
|
|||
|
|
print("警告: 模型没有可训练参数")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"模型验证警告: {e}")
|
|||
|
|
|
|||
|
|
def _predict_with_mlu(self, image) -> dict:
|
|||
|
|
"""在MLU上执行推理"""
|
|||
|
|
try:
|
|||
|
|
start_time = time.perf_counter()
|
|||
|
|
|
|||
|
|
# 预处理
|
|||
|
|
inputs = self.processor(images=image, return_tensors="pt")
|
|||
|
|
|
|||
|
|
if self.use_mlu:
|
|||
|
|
# 将输入数据移动到MLU
|
|||
|
|
inputs_mlu = {}
|
|||
|
|
for key, value in inputs.items():
|
|||
|
|
if hasattr(value, 'mlu'):
|
|||
|
|
inputs_mlu[key] = value.mlu()
|
|||
|
|
else:
|
|||
|
|
inputs_mlu[key] = value
|
|||
|
|
|
|||
|
|
# 执行推理
|
|||
|
|
with torch.no_grad():
|
|||
|
|
# 首次推理(热身)
|
|||
|
|
ts = time.time()
|
|||
|
|
outputs = self.model(**inputs_mlu)
|
|||
|
|
#first_pass_time = time.time() - ts
|
|||
|
|
print('mlu370 T1', time.time() - ts, flush=True)
|
|||
|
|
|
|||
|
|
# 多次推理(性能测试)
|
|||
|
|
ts = time.time()
|
|||
|
|
#for _ in range(5): # 减少测试次数
|
|||
|
|
for i in range(800):
|
|||
|
|
outputs = self.model(**inputs_mlu)
|
|||
|
|
#batch_pass_time = time.time() - ts
|
|||
|
|
print('mlu370 T2', time.time() - ts, flush=True)
|
|||
|
|
else:
|
|||
|
|
# CPU推理
|
|||
|
|
with torch.no_grad():
|
|||
|
|
ts = time.time()
|
|||
|
|
outputs = self.model(**inputs)
|
|||
|
|
#first_pass_time = time.time() - ts
|
|||
|
|
print('cpu T1', time.time() - ts, flush=True)
|
|||
|
|
|
|||
|
|
ts = time.time()
|
|||
|
|
#for _ in range(5):
|
|||
|
|
|
|||
|
|
outputs = self.model(**inputs)
|
|||
|
|
#batch_pass_time = time.time() - ts
|
|||
|
|
print('cpu T2', time.time() - ts, flush=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 计算结果
|
|||
|
|
logits = outputs.logits
|
|||
|
|
probs = torch.nn.functional.softmax(logits, dim=-1)
|
|||
|
|
max_prob, max_idx = probs.max(dim=-1)
|
|||
|
|
class_idx = max_idx.item()
|
|||
|
|
|
|||
|
|
processing_time = round(time.perf_counter() - start_time, 6)
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"class_id": class_idx,
|
|||
|
|
"class_name": self.id2label.get(class_idx, f"class_{class_idx}"),
|
|||
|
|
"confidence": float(max_prob.item()),
|
|||
|
|
"device_used": "mlu" if self.use_mlu else "cpu",
|
|||
|
|
"processing_time": processing_time
|
|||
|
|
|
|||
|
|
}
|
|||
|
|
except Exception as e:
|
|||
|
|
return {
|
|||
|
|
"class_id": -1,
|
|||
|
|
"class_name": "error",
|
|||
|
|
"confidence": 0.0,
|
|||
|
|
"device_used": "mlu" if self.use_mlu else "cpu",
|
|||
|
|
"processing_time": 0.0,
|
|||
|
|
"error": str(e)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def predict(self, image) -> dict:
|
|||
|
|
"""预测入口"""
|
|||
|
|
return self._predict_with_mlu(image)
|
|||
|
|
|
|||
|
|
# 初始化Flask应用
|
|||
|
|
app = Flask(__name__)
|
|||
|
|
|
|||
|
|
# 全局模型加载
|
|||
|
|
try:
|
|||
|
|
MODEL_PATH = os.environ.get("MODEL_PATH", "/model")
|
|||
|
|
print(f"从路径加载模型: {MODEL_PATH}")
|
|||
|
|
classifier = MLUImageClassifier(MODEL_PATH)
|
|||
|
|
print("模型加载成功")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"服务初始化失败: {str(e)}")
|
|||
|
|
classifier = None
|
|||
|
|
|
|||
|
|
@app.route('/v1/private/s782b4996', methods=['POST'])
|
|||
|
|
def predict():
|
|||
|
|
"""接收单张图片并返回GPU预测结果"""
|
|||
|
|
if classifier is None:
|
|||
|
|
return jsonify({
|
|||
|
|
"status": "error",
|
|||
|
|
"prediction": {
|
|||
|
|
"class_id": -1,
|
|||
|
|
"class_name": "error",
|
|||
|
|
"confidence": 0.0,
|
|||
|
|
"device_used": "unknown",
|
|||
|
|
"processing_time": 0.0,
|
|||
|
|
"error": "服务未初始化成功"
|
|||
|
|
}
|
|||
|
|
}), 500
|
|||
|
|
|
|||
|
|
if 'image' not in request.files:
|
|||
|
|
return jsonify({
|
|||
|
|
"status": "error",
|
|||
|
|
"prediction": {
|
|||
|
|
"class_id": -1,
|
|||
|
|
"class_name": "error",
|
|||
|
|
"confidence": 0.0,
|
|||
|
|
"device_used": "mlu" if classifier.use_mlu else "cpu",
|
|||
|
|
"processing_time": 0.0,
|
|||
|
|
"error": "请求中未包含图片"
|
|||
|
|
}
|
|||
|
|
}), 400
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
image_file = request.files['image']
|
|||
|
|
image = Image.open(BytesIO(image_file.read())).convert("RGB")
|
|||
|
|
result = classifier.predict(image)
|
|||
|
|
|
|||
|
|
if 'error' in result:
|
|||
|
|
return jsonify({
|
|||
|
|
"status": "error",
|
|||
|
|
"prediction": result
|
|||
|
|
}), 500
|
|||
|
|
else:
|
|||
|
|
return jsonify({
|
|||
|
|
"status": "success",
|
|||
|
|
"prediction": result
|
|||
|
|
})
|
|||
|
|
except Exception as e:
|
|||
|
|
return jsonify({
|
|||
|
|
"status": "error",
|
|||
|
|
"prediction": {
|
|||
|
|
"class_id": -1,
|
|||
|
|
"class_name": "error",
|
|||
|
|
"confidence": 0.0,
|
|||
|
|
"device_used": "mlu" if classifier and classifier.use_mlu else "cpu",
|
|||
|
|
"processing_time": 0.0,
|
|||
|
|
"error": f"处理图片失败: {str(e)}"
|
|||
|
|
}
|
|||
|
|
}), 500
|
|||
|
|
|
|||
|
|
@app.route('/health', methods=['GET'])
|
|||
|
|
def health_check():
|
|||
|
|
"""健康检查接口"""
|
|||
|
|
mlu_available = False
|
|||
|
|
mlu_info = {}
|
|||
|
|
|
|||
|
|
if torch_mlu is not None and hasattr(ct, 'is_mlu_available'):
|
|||
|
|
try:
|
|||
|
|
mlu_available = ct.is_mlu_available()
|
|||
|
|
mlu_info = {
|
|||
|
|
"device_count": ct.device_count(),
|
|||
|
|
"devices": [ct.get_device_name(i) for i in range(ct.device_count())]
|
|||
|
|
}
|
|||
|
|
except Exception as e:
|
|||
|
|
mlu_info["error"] = str(e)
|
|||
|
|
|
|||
|
|
return jsonify({
|
|||
|
|
"status": "healthy" if classifier is not None else "degraded",
|
|||
|
|
"mlu_available": mlu_available,
|
|||
|
|
"mlu_info": mlu_info,
|
|||
|
|
"model_loaded": classifier is not None,
|
|||
|
|
"using_mlu": classifier.use_mlu if classifier else False,
|
|||
|
|
"timestamp": time.time()
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
@app.route('/test', methods=['GET'])
|
|||
|
|
def test_mlu():
|
|||
|
|
"""MLU测试接口"""
|
|||
|
|
try:
|
|||
|
|
if torch_mlu is None:
|
|||
|
|
return jsonify({
|
|||
|
|
"status": "error",
|
|||
|
|
"message": "torch_mlu模块未找到",
|
|||
|
|
"mlu_working": False
|
|||
|
|
}), 500
|
|||
|
|
|
|||
|
|
# 测试MLU基本功能
|
|||
|
|
test_tensor = torch.randn(3, 3).mlu()
|
|||
|
|
result_tensor = test_tensor + test_tensor
|
|||
|
|
result_cpu = result_tensor.cpu()
|
|||
|
|
|
|||
|
|
return jsonify({
|
|||
|
|
"status": "success",
|
|||
|
|
"message": "MLU测试通过",
|
|||
|
|
"result_shape": str(result_cpu.shape),
|
|||
|
|
"mlu_working": True
|
|||
|
|
})
|
|||
|
|
except Exception as e:
|
|||
|
|
return jsonify({
|
|||
|
|
"status": "error",
|
|||
|
|
"message": f"MLU测试失败: {str(e)}",
|
|||
|
|
"mlu_working": False
|
|||
|
|
}), 500
|
|||
|
|
|
|||
|
|
@app.route('/info', methods=['GET'])
|
|||
|
|
def device_info():
|
|||
|
|
"""设备信息接口"""
|
|||
|
|
info = {
|
|||
|
|
"pytorch_version": torch.__version__,
|
|||
|
|
"torch_mlu_available": torch_mlu is not None,
|
|||
|
|
"mlu_devices_count": ct.device_count() if torch_mlu and hasattr(ct, 'device_count') else 0,
|
|||
|
|
"model_loaded": classifier is not None,
|
|||
|
|
"using_mlu": classifier.use_mlu if classifier else False,
|
|||
|
|
"system_time": time.time()
|
|||
|
|
}
|
|||
|
|
return jsonify(info)
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# 启动HTTP服务 - 使用Flask内置服务器
|
|||
|
|
print("启动MLU图像分类服务...")
|
|||
|
|
app.run(host='0.0.0.0', port=80, debug=False)
|