PyTorch Java 高校计算机硕士研一课程

Spark 4.0 集成 JavaCPP-PyTorch (2.10-1.5.13) 实战：特征工程与模型训练一体化（PyTorch on Java）

在大数据AI落地场景中，特征工程依赖Spark的分布式计算能力处理海量数据，而模型训练常采用PyTorch框架实现高精度建模。如何在Java生态中，用Spark 4.0完成分布式特征工程，同时通过JavaCPP-PyTorch直接运行PyTorch模型训练，避免跨语言调用的繁琐与性能损耗？本文将基于Spark 4.0与JavaCPP-PyTorch 2.10-1.5.13，实现“特征工程+模型训练”全流程Java化，打造高效、可扩展的大数据AI一体化解决方案。

本文将手把手带你完成Spark 4.0分布式特征工程开发、JavaCPP-PyTorch 2.10-1.5.13环境集成、PyTorch模型训练与评估，覆盖环境搭建、依赖配置、代码实现、实战测试全流程，适配生产级大数据AI场景。

一、技术选型与核心概念

1. 核心组件介绍

• Spark 4.0：新一代大数据分布式计算框架，提供强大的DataFrame/Dataset API，支持分布式特征工程（缺失值填充、编码、归一化、特征选择等），兼容Java生态，处理TB级海量数据高效稳定，是企业级特征工程的首选工具。

• JavaCPP-PyTorch 2.10-1.5.13：基于JavaCPP封装的PyTorch Java绑定库，无需安装Python、无需部署PyTorch环境，直接在JVM上调用PyTorch C++核心API，支持模型定义、训练、评估全流程，兼容PyTorch 2.1核心特性，支持CPU/GPU加速，完美适配Java/Spark生态。

• 版本匹配：本次使用Spark 4.0、JavaCPP-PyTorch 2.10-1.5.13（对应PyTorch 2.1核心版本），依赖JDK 17+，兼容Windows/Linux/Mac多平台，支持分布式特征工程与本地/分布式模型训练。

2. 适用场景

• 海量数据场景下，需用Spark做分布式特征工程（如用户行为特征、文本特征、时序特征处理），同时用PyTorch做模型训练（分类、回归、深度学习等）。

• 希望在Java生态中完成“特征工程+模型训练”全流程，避免Python与Java跨语言调用（如Socket、RPC）的性能损耗与部署复杂度。

• 企业级大数据AI项目，要求特征工程可扩展、模型训练可复用，且整体架构基于Java技术栈，降低多语言维护成本。

• 需要利用GPU加速模型训练，同时依托Spark的分布式能力处理海量输入特征。

二、环境准备

1. 开发环境：JDK 17+、Maven 3.8+、Spark 4.0（单机/集群均可）、IntelliJ IDEA（推荐）。

2. 依赖准备：提前配置Spark 4.0环境（环境变量SPARK_HOME），确保spark-shell可正常运行；JavaCPP-PyTorch依赖将通过Maven自动下载，无需手动部署PyTorch。

3. 数据准备：准备用于特征工程的样本数据（如CSV/Parquet格式），包含特征列与标签列（用于模型训练），可选用公开数据集（如Iris、Boston Housing）或业务数据。

4. GPU环境（可选）：若需GPU加速模型训练，需安装NVIDIA CUDA（建议11.8+），JavaCPP-PyTorch将自动识别并调用GPU。

三、Spark 4.0 + JavaCPP-PyTorch 项目搭建

1. 创建Maven项目

初始化Maven项目，引入Spark 4.0核心依赖、JavaCPP-PyTorch依赖，确保项目可正常加载Spark分布式计算能力与PyTorch模型训练API。

2. 核心Maven依赖（关键）

核心依赖包含Spark 4.0核心组件、JavaCPP-PyTorch平台通用依赖，自动适配对应操作系统的底层库，pom.xml配置如下：

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>
    <groupId>com.example</groupId>
    <artifactId>spark4-pytorch-demo</artifactId>
    <version>0.0.1-SNAPSHOT</version>
    <name>spark4-pytorch-demo</name>
    <description>Spark 4.0 集成 JavaCPP-PyTorch 2.10-1.5.13 实战（特征工程+模型训练）</description>

    <properties>
        <java.version>17</java.version>
        <spark.version>4.0.0</spark.version>
        <javacpp.version>1.5.13</javacpp.version>
        <pytorch.version>2.10-1.5.13</pytorch.version>
        <scala.version>2.12</scala.version>
    </properties>

    <dependencies&gt;
        <!-- Spark 4.0 核心依赖（DataFrame API、特征工程） -->
        <dependency>
            <groupId>org.apache.spark</groupId>
            <artifactId>spark-sql_${scala.version}</artifactId>
            <version>${spark.version}</version>
            <scope>provided</scope>
        </dependency>
        <dependency>
            <groupId>org.apache.spark</groupId>
            <artifactId>spark-mllib_${scala.version}</artifactId>
            <version>${spark.version}</version>
            <scope>provided</scope&gt;
        &lt;/dependency&gt;

        <!-- JavaCPP 核心依赖 -->
        <dependency>
            <groupId>org.bytedeco</groupId>
            <artifactId>javacpp</artifactId>
            <version>${javacpp.version}</version>
        &lt;/dependency&gt;

        <!-- PyTorch Java 绑定 平台通用依赖（核心，模型训练） -->
        <dependency>
            <groupId>org.bytedeco</groupId>
            <artifactId>pytorch-platform</artifactId>
            <version>${pytorch.version}&lt;/version&gt;
        &lt;/dependency&gt;

        <!-- 测试依赖 -->
        <dependency>
            <groupId>junit</groupId>
            <artifactId>junit</artifactId>
            <version>4.13.2</version>
            <scope>test</scope>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-compiler-plugin</artifactId>
                <version>3.11.0</version>
                <configuration>
                    <source>${java.version}</source>
                    <target>${java.version}</target>
                    <encoding>UTF-8</encoding>
                </configuration>
            </plugin&gt;
            <!-- 打包插件，适配Spark集群部署 -->
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-shade-plugin</artifactId>
                <version>3.4.1</version>
                <executions>
                    <execution>
                        <phase>package</phase>
                        <goals>
                            <goal>shade</goal>
                        </goals>
                        <configuration>
                            <filters>
                                <filter>
                                    <artifact>*:*</artifact>
                                    <excludes>
                                        <exclude>META-INF/*.SF</exclude>
                                        <exclude>META-INF/*.DSA</exclude>
                                        <exclude>META-INF/*.RSA</exclude>
                                    </excludes>
                                </filter>
                            </filters>
                        </configuration>
                    </execution>
                </executions>
            </plugin>
        </plugins>
    </build>
</project>

注意：首次加载会自动下载600MB+底层库（包含Spark依赖与PyTorch核心运行时），属于正常现象；Spark依赖标注为provided，集群部署时无需打包Spark相关依赖，避免冲突。

四、Spark 4.0 分布式特征工程实现

基于Spark 4.0 DataFrame API，实现海量数据的特征工程全流程，包含数据读取、缺失值处理、特征编码、归一化、特征选择，最终输出适配PyTorch训练的特征张量。

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.ml.feature.*;
import org.apache.spark.ml.linalg.DenseVector;
import org.apache.spark.ml.linalg.Vector;
import java.util.Arrays;

/**
 * Spark 4.0 分布式特征工程工具类
 * 负责数据读取、特征预处理、特征转换，输出适配PyTorch训练的特征数据
 */
public class SparkFeatureEngineering {

    // SparkSession 单例对象（分布式计算核心）
    private SparkSession sparkSession;

    /**
     * 初始化SparkSession（适配Spark 4.0）
     */
    public void initSpark() {
        sparkSession = SparkSession.builder()
                .appName("Spark4FeatureEngineering")
                .master("local[*]") // 本地模式，集群部署时删除该配置
                .config("spark.driver.memory", "4g")
                .config("spark.executor.memory", "8g")
                .getOrCreate();
        System.out.println("✅ Spark 4.0 初始化成功，Spark版本：" + sparkSession.version());
    }

    /**
     * 1. 读取数据（支持CSV/Parquet格式）
     * @param dataPath 数据路径（本地/分布式文件系统）
     * @return 原始DataFrame
     */
    public Dataset<Row> readData(String dataPath) {
        return sparkSession.read()
                .option("header", "true")
                .option("inferSchema", "true")
                .csv(dataPath); // 若为Parquet格式，替换为.parquet(dataPath)
    }

    /**
     * 2. 特征工程全流程（核心方法）
     * @param rawData 原始DataFrame（包含特征列、标签列）
     * @param featureCols 特征列名数组
     * @param labelCol 标签列名
     * @return 处理后的DataFrame（包含特征向量、标签）
     */
    public Dataset<Row> processFeatures(Dataset<Row> rawData, String[] featureCols, String labelCol) {
        // 步骤1：缺失值填充（数值型特征用均值填充）
        Imputer imputer = new Imputer()
                .setInputCols(featureCols)
                .setOutputCols(Arrays.stream(featureCols).map(col -> col + "_imputed").toArray(String[]::new))
                .setStrategy("mean");
        Dataset<Row> imputedData = imputer.fit(rawData).transform(rawData);

        // 步骤2：特征归一化（标准化，适配PyTorch模型训练）
        StandardScaler scaler = new StandardScaler()
                .setInputCol("features")
                .setOutputCol("scaled_features")
                .setWithMean(true)
                .setWithStd(true);

        // 步骤3：将特征列合并为特征向量
        VectorAssembler assembler = new VectorAssembler()
                .setInputCols(Arrays.stream(featureCols).map(col -> col + "_imputed").toArray(String[]::new))
                .setOutputCol("features");
        Dataset<Row> featureData = assembler.transform(imputedData);

        // 步骤4：特征归一化
        Dataset<Row> scaledData = scaler.fit(featureData).transform(featureData);

        // 步骤5：选择最终列（归一化后的特征向量、标签）
        return scaledData.select("scaled_features", labelCol)
                .withColumnRenamed(labelCol, "label");
    }

    /**
     * 3. 特征向量转float数组（适配PyTorch张量）
     * @param featureVector Spark特征向量
     * @return 对应float数组
     */
    public float[] vectorToFloatArray(Vector featureVector) {
        DenseVector denseVector = (DenseVector) featureVector;
        return denseVector.values();
    }

    /**
     * 关闭SparkSession，释放资源
     */
    public void closeSpark() {
        if (sparkSession != null) {
            sparkSession.stop();
            System.out.println("✅ Spark 4.0 资源已释放");
        }
    }
}

五、JavaCPP-PyTorch 2.10-1.5.13 模型训练封装

创建PyTorch模型训练工具类，基于JavaCPP-PyTorch API定义模型结构、实现训练逻辑、评估模型性能，接收Spark特征工程输出的特征数据，完成模型训练与保存，适配Java生态调用。

import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.global.torch;
import java.io.File;
import java.util.ArrayList;
import java.util.List;

/**
 * JavaCPP-PyTorch 2.10-1.5.13 模型训练工具类
 * 定义PyTorch模型、实现训练与评估，接收Spark特征数据，适配Java调用
 */
public class PyTorchTrainManager {

    // PyTorch模型对象
    private Module model;
    // 优化器
    private Optimizer optimizer;
    // 损失函数（以回归任务为例，分类任务可替换为CrossEntropyLoss）
    private Loss criterion;
    // 模型保存路径
    private final String modelSavePath;

    public PyTorchTrainManager(String modelSavePath) {
        this.modelSavePath = modelSavePath;
        // 初始化模型、优化器、损失函数
        initModel();
    }

    /**
     * 初始化PyTorch模型（线性回归模型，可根据需求替换为深度学习模型）
     */
    private void initModel() {
        // 定义模型：输入维度（根据特征数量调整）、输出维度（回归为1，分类为类别数）
        int inputDim = 4; // 示例：4个特征
        int outputDim = 1;

        // 构建线性模型（Sequential）
        model = new Sequential(
                new Linear(inputDim, 16), // 隐藏层1：输入inputDim，输出16
                new ReLU(), // 激活函数
                new Linear(16, 8), // 隐藏层2：输入16，输出8
                new ReLU(), // 激活函数
                new Linear(8, outputDim) // 输出层
        );

        // 定义优化器（Adam）
        optimizer = torch.optim_adam(model.parameters(), new AdamOptions(0.001));

        // 定义损失函数（MSELoss，适用于回归任务）
        criterion = new MSELoss();

        // 打印模型结构
        System.out.println("✅ PyTorch模型初始化完成，模型结构：");
        System.out.println(model);

        // 检查GPU是否可用
        if (torch.cuda_is_available()) {
            model.to(torch.device(torch.CUDA));
            criterion.to(torch.device(torch.CUDA));
            System.out.println("✅ GPU可用，已切换至GPU训练");
        } else {
            System.out.println("ℹ️ GPU不可用，使用CPU训练");
        }
    }

    /**
     * 模型训练核心方法
     * @param features 特征数组列表（来自Spark特征工程）
     * @param labels 标签数组列表
     * @param epochs 训练轮次
     * @param batchSize 批次大小
     */
    public void train(List<float[]> features, List<Float> labels, int epochs, int batchSize) {
        // 将特征和标签转换为PyTorch张量
        Tensor featureTensor = createFeatureTensor(features);
        Tensor labelTensor = createLabelTensor(labels);

        // 开启训练模式
        model.train(true);

        for (int epoch = 0; epoch < epochs; epoch++) {
            // 清零梯度
            optimizer.zero_grad();

            // 前向传播
            Tensor outputs = model.forward(featureTensor).toTensor();

            // 计算损失
            Tensor loss = criterion.forward(outputs, labelTensor);

            // 反向传播、参数更新
            loss.backward();
            optimizer.step();

            // 打印训练日志
            if ((epoch + 1) % 10 == 0) {
                float lossValue = loss.item().get_float();
                System.out.printf("Epoch [%d/%d], Loss: %.4f%n", epoch + 1, epochs, lossValue);
            }
        }

        // 训练完成，保存模型（TorchScript格式，可后续推理使用）
        saveModel();

        // 释放张量资源
        featureTensor.close();
        labelTensor.close();
        System.out.println("✅ 模型训练完成，已保存至：" + modelSavePath);
    }

    /**
     * 特征列表转PyTorch张量
     */
    private Tensor createFeatureTensor(List<float[]> features) {
        int batchSize = features.size();
        int featureDim = features.get(0).length;
        float[] data = new float[batchSize * featureDim];

        // 拼接所有特征数组
        for (int i = 0; i < batchSize; i++) {
            System.arraycopy(features.get(i), 0, data, i * featureDim, featureDim);
        }

        // 创建张量（形状：[batchSize, featureDim]）
        return torch.tensor_from_float_buffer(data, new long[]{batchSize, featureDim})
                .to(torch.get_default_dtype());
    }

    /**
     * 标签列表转PyTorch张量
     */
    private Tensor createLabelTensor(List<Float> labels) {
        float[] data = new float[labels.size()];
        for (int i = 0; i < labels.size(); i++) {
            data[i] = labels.get(i);
        }
        // 创建张量（形状：[batchSize, 1]）
        return torch.tensor_from_float_buffer(data, new long[]{labels.size(), 1})
                .to(torch.get_default_dtype());
    }

    /**
     * 保存TorchScript模型
     */
    private void saveModel() {
        File modelDir = new File(modelSavePath).getParentFile();
        if (!modelDir.exists()) {
            modelDir.mkdirs();
        }
        // 导出为TorchScript模型
        ScriptModule scriptModule = torch.jit_script(model);
        scriptModule.save(modelSavePath);
        scriptModule.close();
    }

    /**
     * 模型评估（计算预测误差）
     * @param features 测试特征
     * @param labels 测试标签
     * @return 平均误差
     */
    public float evaluate(List<float[]> features, List<Float> labels) {
        model.eval(); // 开启评估模式
        Tensor featureTensor = createFeatureTensor(features);
        Tensor labelTensor = createLabelTensor(labels);

        // 前向传播（禁用梯度计算，提升性能）
        try (NoGradGuard guard = new NoGradGuard()) {
            Tensor outputs = model.forward(featureTensor).toTensor();
            Tensor loss = criterion.forward(outputs, labelTensor);
            float avgLoss = loss.item().get_float();

            // 释放资源
            featureTensor.close();
            labelTensor.close();
            return avgLoss;
        }
    }

    /**
     * 释放模型资源
     */
    public void close() {
        if (model != null) {
            model.close();
        }
        if (optimizer != null) {
            optimizer.close();
        }
        if (criterion != null) {
            criterion.close();
        }
        System.out.println("✅ PyTorch模型资源已释放");
    }
}

六、全流程整合与实战测试

整合Spark 4.0特征工程与JavaCPP-PyTorch模型训练，实现“数据读取→特征预处理→模型训练→评估”全流程，以Iris数据集（分类任务）或Boston Housing数据集（回归任务）为例，完成实战测试。

1. 全流程主程序

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import java.util.ArrayList;
import java.util.List;

/**
 * 主程序：Spark 4.0 特征工程 + JavaCPP-PyTorch 2.10-1.5.13 模型训练全流程
 */
public class SparkPyTorchFullProcess {
    public static void main(String[] args) {
        // 1. 初始化Spark特征工程工具
        SparkFeatureEngineering featureEngineering = new SparkFeatureEngineering();
        featureEngineering.initSpark();

        try {
            // 2. 读取数据（示例：Iris数据集，本地路径）
            String dataPath = "src/main/resources/iris.csv";
            Dataset<Row> rawData = featureEngineering.readData(dataPath);
            rawData.show(5); // 打印前5条数据

            // 3. 定义特征列和标签列（Iris数据集：4个特征，1个标签）
            String[] featureCols = {"sepal_length", "sepal_width", "petal_length", "petal_width"};
            String labelCol = "species"; // 分类标签（需提前编码，此处简化处理）

            // 4. 执行特征工程
            Dataset<Row> processedData = featureEngineering.processFeatures(rawData, featureCols, labelCol);
            processedData.show(5);

            // 5. 提取特征和标签，转换为PyTorch可接收的格式
            List<float[]> features = new ArrayList<>();
            List<Float> labels = new ArrayList<>();

            // 遍历处理后的DataFrame，提取特征向量和标签
            processedData.foreach(row -> {
                float[] featureArr = featureEngineering.vectorToFloatArray(row.getAs("scaled_features"));
                float label = row.getAs("label"); // 标签需提前转换为数值型
                features.add(featureArr);
                labels.add(label);
            });

            // 6. 初始化PyTorch训练工具
            String modelSavePath = "src/main/resources/pytorch-iris-model.pt";
            PyTorchTrainManager trainManager = new PyTorchTrainManager(modelSavePath);

            // 7. 模型训练（设置训练参数）
            int epochs = 100;
            int batchSize = 32;
            trainManager.train(features, labels, epochs, batchSize);

            // 8. 模型评估（此处用训练集评估，实际场景需拆分训练集/测试集）
            float avgLoss = trainManager.evaluate(features, labels);
            System.out.printf("✅ 模型评估完成，平均损失：%.4f%n", avgLoss);

            // 9. 释放资源
            trainManager.close();
        } catch (Exception e) {
            System.err.println("❌ 全流程执行失败：" + e.getMessage());
            e.printStackTrace();
        } finally {
            // 关闭Spark资源
            featureEngineering.closeSpark();
        }
    }
}

2. 数据准备与测试说明

• 数据集准备：下载Iris数据集（或Boston Housing数据集），保存为CSV格式，放置在src/main/resources目录下，CSV需包含特征列和标签列（标签需为数值型，分类任务可通过Spark的StringIndexer编码）。

• 运行方式：本地运行主程序，Spark将以本地模式执行特征工程，PyTorch自动选择CPU/GPU进行模型训练；集群部署时，修改SparkSession的master配置为集群地址（如yarn），并打包为jar包提交。

• 预期结果：程序将输出Spark特征工程结果、PyTorch模型结构、训练日志、评估损失，最终在指定路径生成TorchScript模型文件，完成全流程验证。

七、关键优化与生产注意事项

1. 性能优化

• Spark特征工程优化：使用Spark 4.0新特性（如向量ized操作、分区优化），减少Shuffle操作；特征工程过程中可启用缓存（data.cache()），提升重复计算效率。

• PyTorch训练优化：开启model.eval()禁用梯度计算（评估时）；使用批量训练（batchSize根据内存调整）；GPU环境下，确保张量与模型均切换至GPU，提升训练速度3-10倍。

• 数据传输优化：Spark特征数据转PyTorch张量时，尽量批量处理，避免单条数据循环转换，减少内存开销。

2. 分布式适配

• Spark 4.0支持集群部署，特征工程可分布式处理海量数据，无需修改代码，仅需调整SparkSession的master配置（yarn/local/k8s）。

• PyTorch模型训练支持单机多GPU、分布式训练（需额外配置），可根据数据量和算力需求调整，适配生产级大规模训练场景。

3. 模型与数据管理

• 模型保存为TorchScript格式，可后续用于JavaCPP-PyTorch推理，也可导出至Python端进行后续优化。

• Spark特征工程输出的数据可保存为Parquet格式，便于后续复用；训练过程中建议拆分训练集、测试集、验证集，提升模型泛化能力。

• 内存管理：及时释放Spark DataFrame、PyTorch张量资源，避免JVM内存溢出；集群部署时，合理配置Spark driver/executor内存。

4. 版本兼容注意

• 确保Spark 4.0与JDK 17+兼容，避免版本过低导致API报错。

• JavaCPP-PyTorch 2.10-1.5.13对应PyTorch 2.1核心版本，若需使用PyTorch更高版本，需同步升级JavaCPP-PyTorch版本，确保API兼容。

八、总结

通过Spark 4.0 + JavaCPP-PyTorch 2.10-1.5.13，我们实现了“分布式特征工程+PyTorch模型训练”全流程Java化，核心优势如下：

1. 生态无缝融合：Spark的分布式特征工程能力与PyTorch的高精度建模能力完美结合，无需跨语言调用，降低部署与维护成本。

2. 全流程Java化：从数据读取、特征预处理到模型训练、评估，全程基于Java技术栈，适配企业级Java生态需求。

3. 高效可扩展：Spark 4.0处理海量特征数据高效稳定，JavaCPP-PyTorch支持GPU加速，可适配从小规模测试到大规模生产的各类场景。

4. 易用性强：JavaCPP-PyTorch无需部署Python环境，底层库自动加载；Spark 4.0 API简洁易用，特征工程代码可复用性高。

该方案完美解决了大数据场景下“特征工程与模型训练分离”“跨语言调用繁琐”的痛点，适用于推荐系统、风控建模、图像/文本大数据建模等各类企业级AI场景，是大数据与AI一体化落地的最优方案之一。