Data+AI 场景实践教程

import ray
import time
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("RayTrainer")

# 清理已有Ray连接
logger.info("清理已有Ray连接...")
if ray.is_initialized():
    ray.shutdown()
    time.sleep(1)  
    logger.info("成功关闭历史Ray连接")

# 初始化Ray连接
logger.info("初始化新的Ray连接...")
ray.init(address="auto", ignore_reinit_error=True)
logger.info("Ray自动连接成功！")

import raydp
from pyspark.sql import SparkSession

# 初始化RayDP Spark会话
spark = raydp.init_spark(
    app_name="MNIST_Preprocessing",
    num_executors=2,
    executor_cores=2,
    executor_memory="1GB",  # 需根据集群情况配置执行节点数、内存
    configs={
        "spark.rpc.message.maxSize": "512",
        "spark.driver.maxResultSize": "512m"
    }
)
logger.info("Spark会话创建成功！")

import torchvision
import torchvision.transforms as transforms
import pandas as pd
import numpy as np
from pyspark.ml.feature import StandardScaler
from pyspark.sql.types import StructType, StructField, ArrayType, FloatType, IntegerType
from pyspark.sql.functions import udf
from pyspark.ml.linalg import Vectors, VectorUDT

def prepare_dataset_with_spark():
    """使用Spark进行数据预处理"""
    logger.info("下载MNIST数据集...")
    transform = transforms.ToTensor()
    train_set = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    test_set = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
    
    logger.info("将数据转换为Spark DataFrame...")
    train_images = [x.view(-1).numpy().tolist() for x, _ in train_set]
    train_labels = [y for _, y in train_set]
    train_df = pd.DataFrame({
        "features": train_images,
        "label": train_labels
    })
    
    test_images = [x.view(-1).numpy().tolist() for x, _ in test_set]
    test_labels = [y for _, y in test_set]
    test_df = pd.DataFrame({
        "features": test_images,
        "label": test_labels
    })
    
    logger.info("创建Spark DataFrame...")
    schema = StructType([
        StructField("features", ArrayType(FloatType()), True),
        StructField("label", IntegerType(), True)
    ])
    
    spark_train_df = spark.createDataFrame(train_df, schema=schema)
    spark_test_df = spark.createDataFrame(test_df, schema=schema)
    
    logger.info("使用Spark进行标准化...")
    array_to_vector = udf(lambda arr: Vectors.dense(arr), VectorUDT())
    
    spark_train_df = spark_train_df.withColumn("features_vec", array_to_vector("features"))
    spark_test_df = spark_test_df.withColumn("features_vec", array_to_vector("features"))
    
    scaler = StandardScaler(inputCol="features_vec", outputCol="scaled_features", withStd=True, withMean=True)
    scaler_model = scaler.fit(spark_train_df)
    spark_train_df = scaler_model.transform(spark_train_df)
    spark_test_df = scaler_model.transform(spark_test_df)
    
    return spark_train_df, spark_test_df

from torch.utils.data import DataLoader, Subset, TensorDataset

def spark_df_to_torch_dataset(spark_df):
    """将Spark DataFrame转换为PyTorch Dataset"""
    pandas_df = spark_df.select("scaled_features", "label").toPandas()
    features = np.stack(pandas_df['scaled_features'].apply(lambda x: x.toArray()))
    labels = pandas_df['label'].values
    
    features_tensor = torch.tensor(features, dtype=torch.float32)
    labels_tensor = torch.tensor(labels, dtype=torch.long)
    
    return TensorDataset(features_tensor, labels_tensor)

@ray.remote
def train_worker(worker_id, num_workers, train_dataset, test_dataset, epochs=3):
    """每个worker的训练任务"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 分配数据给每个worker
    worker_indices = list(range(worker_id, len(train_dataset), num_workers))
    worker_train_set = Subset(train_dataset, worker_indices)
    
    train_loader = DataLoader(worker_train_set, batch_size=64, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
    
    model = SimpleCNN().to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    
    # 训练过程
    for epoch in range(epochs):
        model.train()
        total_loss = 0
        correct = 0
        total = 0
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
        
        train_accuracy = 100 * correct / total
        model.eval()
        test_correct, test_total = 0, 0
        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                _, predicted = torch.max(output.data, 1)
                test_total += target.size(0)
                test_correct += (predicted == target).sum().item()
        
        test_accuracy = 100 * test_correct / test_total
        return worker_id, test_accuracy

import mlflow
import mlflow.pytorch

def log_model_to_mlflow(model_path, accuracy):
    """将模型记录到MLflow"""
    mlflow.set_tracking_uri("<http://localhost:5000")>  # 修改为跟踪服务器
    mlflow.set_experiment("MNIST_Classification")
    
    with mlflow.start_run():
        mlflow.log_param("model_type", "SimpleCNN")
        mlflow.log_param("optimizer", "Adam")
        mlflow.log_param("learning_rate", 0.001)
        mlflow.log_metric("accuracy", accuracy)
        
        model = SimpleCNN()
        model.load_state_dict(torch.load(model_path))
        mlflow.pytorch.log_model(model, "model", registered_model_name="MNIST_Classifier")
        mlflow.log_artifact(model_path, "model_files")

def offline_inference(model, data):
    """使用训练好的模型进行推理"""
    model.eval()
    with torch.no_grad():
        inputs = torch.tensor(data, dtype=torch.float32)
        predictions = model(inputs)
        # 获取每个样本的预测类别（即最大概率对应的类别）
        predicted_classes = torch.argmax(predictions, dim=1)
        return predicted_classes

def perform_offline_inference(model_path, data):
    """执行离线推理并保存结果"""
    logger.info(f"加载模型: {model_path}...")
    
    # 加载训练好的模型
    model = SimpleCNN()
    model.load_state_dict(torch.load(model_path))
    
    # 执行离线推理
    logger.info("开始离线推理...")
    predictions = offline_inference(model, data)
    
    # 将预测类别保存到CSV文件
    result_df = pd.DataFrame(predictions.numpy(), columns=["predictions"])
    result_df.to_csv("offline_predictions.csv", index=False)
    
    logger.info("离线推理完成，结果已保存至 offline_predictions.csv")