Feature Request: Add Support for torch.jit.script, torch.jit.trace, and AOT Optimization APIs in JavaCPP PyTorch Background

[mullerhai]

Hi @saudet
I am using JavaCPP PyTorch for PyTorch model deployment in Java applications. My core requirement is to perform PyTorch AOT (Ahead-of-Time) compilation directly in Java to generate lightweight, Python-independent inference models (optimized TorchScript modules).

      Currently, JavaCPP PyTorch lacks critical JIT compilation and AOT optimization APIs that are natively available in LibTorch C++. This blocks full AOT model conversion workflows in Java.

Missing Core APIs

I need the following torch::jit functions implemented in JavaCPP:
torch::jit::script(...)
torch::jit::trace(...) not like torch.trace is math
torch::jit::load(...)
torch::jit::freeze(Module module) // we have implement
torch::jit::optimize_for_inference(Module module) // we have implement
Full binding for torch::jit::save(Module module, const std::string& path)
Use Case
I want to implement TorchScript → AOT optimized model conversion in Java, identical to this working LibTorch C++ logic:

运行
// C++ Reference (Working)
torch::jit::Module module = torch::jit::load(input_pt_path);
module.eval();
auto frozen = torch::jit::freeze(module);
torch::jit::optimize_for_inference(frozen);
torch::jit::save(frozen, output_aot_path);

Expected Java Implementation (JavaCPP)
This is the Java code I expect to write with the added APIs:

运行
import org.bytedeco.pytorch.*;
import static org.bytedeco.pytorch.global.torch.*;
import static org.bytedeco.pytorch.global.torch_jit.*;

public class TorchAotConverter {
    public static void convertToAOT(String inputPtPath, String outputAotPath) {
        // 1. Load TorchScript model
        Module module = torch.load(inputPtPath);
        
        // 2. Switch to evaluation mode
        module.eval();
        
        // 3. Freeze module (AOT required)
        Module frozenModule = torch.freeze(module);
        
        // 4. Run AOT inference optimization
        torch.optimize_for_inference(frozenModule);
        
        // 5. Save final AOT model
        torch_jit.save(frozenModule, outputAotPath); // we lack it
    }

    public static void main(String[] args) {
        convertToAOT("model.pt", "model.aot");
    }
}

```
Value & Impact
Enables full AOT model compilation in Java (no Python dependency)
Matches production-grade LibTorch C++ deployment capabilities
Critical for high-performance Java + PyTorch inference deployment
Unlocks model optimization, freezing, and porting workflows for Java developers
I would greatly appreciate adding these essential torch::jit scripting/tracing/AOT APIs to JavaCPP PyTorch. Thank you!
the libtorch Aot compile model
```cpp
#include <torch/torch.h>
#include <torch/script.h>
#include <iostream>
#include <string>

int main(int argc, const char* argv[]) {
    // 1. 检查命令行参数
    if (argc != 3) {
        std::cerr << "用法: " << argv[0] << " <输入.pt模型路径> <输出.aot模型路径>" << std::endl;
        std::cerr << "示例: " << argv[0] << " model.pt model.aot" << std::endl;
        return -1;
    }

    std::string input_pt_path = argv[1];  // 输入的.pt模型
    std::string output_aot_path = argv[2]; // 输出的.aot模型

    try {
        // 2. 加载 TorchScript 格式的 .pt 模型
        std::cout << "正在加载 .pt 模型: " << input_pt_path << std::endl;
        torch::jit::Module module;
        module = torch::jit::load(input_pt_path);

        // 3. 【核心】AOT 编译（启用优化，生成无Python依赖的模型）
        std::cout << "正在执行 AOT 编译优化..." << std::endl;
        // 冻结模型（固定权重，禁止修改，AOT必备步骤）
        module.eval(); // 切换到推理模式
        auto frozen_module = torch::jit::freeze(module);
        // 启用 AOT 编译优化（核心API）
        torch::jit::optimize_for_inference(frozen_module);

        // 4. 保存编译后的 AOT 模型
        std::cout << "正在保存 AOT 模型到: " << output_aot_path << std::endl;
        torch::jit::save(frozen_module, output_aot_path);

        std::cout << "✅ 模型转换完成！.pt → .aot 成功" << std::endl;
        return 0;
    } catch (const c10::Error& e) {
        std::cerr << "❌ 错误: " << e.what() << std::endl;
        return -1;
    }
}

```

[mullerhai]

package torch;

import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.global.torch;
import org.bytedeco.javacpp.PointerScope;
import java.nio.file.*;
import java.util.*;

/**
 * JitAPITest - JavaCPP-PyTorch JIT ?????
 *
 * ????: ?? -> ?? -> ?? -> freeze -> optimize -> ??
 */
public class JitAPITest {

    static class Result {
        String name; boolean passed; long durationMs; String detail;
        Result(String n, boolean p, long d, String det) {
            name = n; passed = p; durationMs = d; detail = det;
        }
    }

    static final String MODEL_PATH = "/tmp/jit_test_model.pt";
    static final String TRAINED_MODEL_PATH = "/tmp/jit_trained_model.pt";

    public static void main(String[] args) throws Exception {
        System.out.println("?????????????????????????????????????????????????????????????????");
        System.out.println("?        JavaCPP-PyTorch JIT ?????                      ?");
        System.out.println("?????????????????????????????????????????????????????????????????");
        System.out.println();

        List<Result> results = new ArrayList<>();

        // Step 1: ????
        results.add(testTrainModel());

        // Step 2: ????
        results.add(testSaveModel());

        // Step 3: Load JIT module from file
        results.add(testLoadJitModuleFromFile());

        // Step 4: Freeze loaded module
        results.add(testFreezeLoadedModule());

        // Step 5: Optimize frozen module
        results.add(testOptimizeFrozenModule());

        // Step 6: ??????
        results.add(testInferenceComparison());

        // Step 7: ?????????
        results.add(testSaveLoadStability());

        // Step 8: ???freeze/optimize
        results.add(testLargeModelFreeze());

        // Print results
        System.out.println("\n???????????????????????????????????????????????????????????????");
        System.out.println("BENCHMARK RESULTS");
        System.out.println("???????????????????????????????????????????????????????????????");

        int passed = 0, failed = 0;
        for (Result r : results) {
            String s = r.passed ? "? PASS" : "? FAIL";
            System.out.printf("%-45s %s (%d ms)%n", r.name, s, r.durationMs);
            if (r.detail != null && !r.detail.isEmpty()) {
                System.out.println("  Detail: " + r.detail);
            }
            if (r.passed) passed++; else failed++;
        }

        System.out.println("\n???????????????????????????????????????????????????????????????");
        System.out.printf("TOTAL: %d/%d PASSED%n", passed, passed + failed);
        System.out.println("???????????????????????????????????????????????????????????????");

        // Cleanup
        Files.deleteIfExists(Paths.get(MODEL_PATH));
        Files.deleteIfExists(Paths.get(TRAINED_MODEL_PATH));

        System.exit(failed > 0 ? 1 : 0);
    }

    // Step 1: ????????
    static Result testTrainModel() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            // ????
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(784, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 10));
            model.eval();

            // ?????
            Optimizer optimizer = new Adam(model.parameters(), new AdamOptions(1e-3));

            // ??????
            int epochs = 5;
            for (int epoch = 0; epoch < epochs; epoch++) {
                // ???????
                Tensor input = torch.randn(new long[]{32, 784});
                Tensor target = torch.randint(10, new long[]{32});

                Tensor output = model.forward(input);
                Tensor loss = torch.cross_entropy(output, target);

                optimizer.zero_grad();
                loss.backward();
                optimizer.step();

                loss.close(); output.close(); target.close(); input.close();
            }

            // ????????
            torch.save(model, TRAINED_MODEL_PATH);

            ok = true;
            detail = "trained " + epochs + " epochs, saved to: " + TRAINED_MODEL_PATH;

            optimizer.close();
            model.close();
        } catch (Exception e) {
            detail = e.getMessage();
            e.printStackTrace();
        }
        return new Result("1. Train Model", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // Step 2: ????
    static Result testSaveModel() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(784, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 10));
            model.eval();

            torch.save(model, MODEL_PATH);

            ok = Files.exists(Paths.get(MODEL_PATH));
            long size = ok ? Files.size(Paths.get(MODEL_PATH)) : 0;
            detail = "model saved, size: " + (size / 1024) + " KB";

            model.close();
        } catch (Exception e) { detail = e.getMessage(); }

        return new Result("2. Save Model", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // Step 3: Load JIT module from file
    static Result testLoadJitModuleFromFile() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            JitModule loaded = torch.load_jit_module_from_file(MODEL_PATH, new ExtraFilesMap());
            ok = loaded != null;
            detail = "loaded type: JitModule";
            loaded.close();
        } catch (Exception e) {
            // ????? ExtraFilesMap
            try {
                JitModule loaded = torch.load_jit_module_from_file(MODEL_PATH);
                ok = loaded != null;
                detail = "loaded type: JitModule (no extra files)";
                loaded.close();
            } catch (Exception e2) {
                detail = "error: " + e2.getMessage();
            }
        }

        return new Result("3. load_jit_module_from_file()", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // Step 4: Freeze loaded module
    static Result testFreezeLoadedModule() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            JitModule loaded = torch.load_jit_module_from_file(MODEL_PATH, new ExtraFilesMap());

            // Freeze
            JitModule frozen = torch.freeze(loaded);
            ok = frozen != null;
            detail = "frozen type: " + (frozen != null ? frozen.getClass().getSimpleName() : "null");

            // ????
            Tensor input = torch.randn(new long[]{1, 784});
            // JitModule ? forward ?? IValueVector
            IValueVector inputs = new IValueVector();
            inputs.push_back(new IValue(input));
            IValue output = frozen.forward(inputs);
            ok = ok && output != null;

            output.close(); input.close(); frozen.close(); loaded.close();
        } catch (Exception e) { detail = e.getMessage(); }

        return new Result("4. torch.freeze(JitModule)", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // Step 5: Optimize frozen module
    static Result testOptimizeFrozenModule() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            JitModule loaded = torch.load_jit_module_from_file(MODEL_PATH, new ExtraFilesMap());
            JitModule frozen = torch.freeze(loaded);

            // Optimize for inference
            JitModule optimized = torch.optimize_for_inference(frozen);
            ok = optimized != null;
            detail = "optimized type: " + (optimized != null ? optimized.getClass().getSimpleName() : "null");

            // ????
            Tensor input = torch.randn(new long[]{1, 784});
            IValueVector inputs = new IValueVector();
            inputs.push_back(new IValue(input));
            IValue output = optimized.forward(inputs);
            ok = ok && output != null;

            output.close(); input.close(); optimized.close(); frozen.close(); loaded.close();
        } catch (Exception e) { detail = e.getMessage(); }

        return new Result("5. torch.optimize_for_inference()", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // Step 6: ??????
    static Result testInferenceComparison() {
        long start = System.nanoTime();
        boolean ok = false;
        int iters = 100;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            // ??????
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(784, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 10));
            model.eval();

            // ??
            for (int i = 0; i < 10; i++) {
                Tensor in = torch.randn(new long[]{1, 784});
                model.forward(in).close();
                in.close();
            }

            // Benchmark ????
            long total1 = 0;
            for (int i = 0; i < iters; i++) {
                Tensor in = torch.randn(new long[]{1, 784});
                long t = System.nanoTime();
                Tensor out = model.forward(in);
                total1 += (System.nanoTime() - t);
                out.close(); in.close();
            }

            // ????? JitModule
            try {
                JitModule loaded = torch.load_jit_module_from_file(MODEL_PATH, new ExtraFilesMap());
                JitModule opt = torch.optimize_for_inference(torch.freeze(loaded));

                IValueVector inputs = new IValueVector();

                long total2 = 0;
                for (int i = 0; i < iters; i++) {
                    Tensor in = torch.randn(new long[]{1, 784});
                    inputs.clear();
                    inputs.push_back(new IValue(in));
                    long t = System.nanoTime();
                    IValue out = opt.forward(inputs);
                    total2 += (System.nanoTime() - t);
                    out.close(); in.close();
                }

                double avg1 = total1 / 1_000_000.0 / iters;
                double avg2 = total2 / 1_000_000.0 / iters;
                double speedup = avg1 / avg2;

                ok = speedup > 0;
                detail = String.format("SEQ: %.2fms, JIT: %.2fms, speedup: %.2fx", avg1, avg2, speedup);

                opt.close(); loaded.close();
            } catch (Exception e) {
                double avg1 = total1 / 1_000_000.0 / iters;
                ok = true;
                detail = "SEQ only: " + avg1 + " ms/iter";
            }

            model.close();
        } catch (Exception e) { detail = e.getMessage(); }

        return new Result("6. Inference Comparison", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // Step 7: ?????????
    static Result testSaveLoadStability() {
        long start = System.nanoTime();
        boolean ok = true;
        String detail = "";
        int cycles = 5;

        try (PointerScope scope = new PointerScope()) {
            for (int i = 0; i < cycles; i++) {
                String path = "/tmp/stable_test_" + i + ".pt";

                SequentialImpl model = new SequentialImpl();
                model.push_back(new LinearImpl(784, 128));
                model.push_back(new ReLUImpl());
                model.push_back(new LinearImpl(128, 10));
                model.eval();

                torch.save(model, path);

                JitModule loaded = torch.load_jit_module_from_file(path, new ExtraFilesMap());
                if (loaded == null) ok = false;

                loaded.close();
                model.close();
                Files.deleteIfExists(Paths.get(path));
            }
            detail = "cycles: " + cycles + " OK";
        } catch (Exception e) {
            ok = false;
            detail = e.getMessage();
        }

        return new Result("7. Save/Load Stability " + cycles + "x", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // Step 8: ??? freeze/optimize
    static Result testLargeModelFreeze() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            // ????????
            SequentialImpl bigModel = new SequentialImpl();
            bigModel.push_back(new LinearImpl(1024, 2048));
            bigModel.push_back(new ReLUImpl());
            bigModel.push_back(new LinearImpl(2048, 2048));
            bigModel.push_back(new ReLUImpl());
            bigModel.push_back(new LinearImpl(2048, 512));
            bigModel.eval();

            String path = "/tmp/large_model.pt";
            torch.save(bigModel, path);
            bigModel.close();

            // Load and freeze
            JitModule loaded = torch.load_jit_module_from_file(path, new ExtraFilesMap());
            JitModule frozen = torch.freeze(loaded);

            // Optimize
            JitModule optimized = torch.optimize_for_inference(frozen);

            // ????
            Tensor input = torch.randn(new long[]{4, 1024});
            IValueVector inputs = new IValueVector();
            inputs.add(new IValue(input));
            IValue output = optimized.forward(inputs);

            ok = !output.isNull();
            detail = "large model (1024->2048->2048->512) processed";

            output.close(); input.close();
            optimized.close(); frozen.close(); loaded.close();
            Files.deleteIfExists(Paths.get(path));
        } catch (Exception e) { detail = e.getMessage(); }

        return new Result("8. Large Model Freeze/Optimize", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }
}

we are lack torch.save(jitmodule) --> aot model file @saudet

[saudet]

Looks like we'd need to add #include <torch/csrc/jit/serialization/export.h> to torch_include.h.

@HGuillemet Any reason in particular why it's not there already?

[mullerhai]

<torch/csrc/jit/serialization/export.h>

Yes ，just this <torch/csrc/jit/serialization/export.h> file ，we need it ，I think maybe forget to compile it on the pre versions.

[mullerhai]

If we not compile torch.jit.save ,we could not export pt format model to aot format ,like AOTPackageLoader , now only could aot loader , like torch.jit.trace and torch.jit.script also important

[mullerhai]

Looks like we'd need to add #include <torch/csrc/jit/serialization/export.h> to torch_include.h.

@HGuillemet Any reason in particular why it's not there already?

If we not compile torch.jit.save ,we could not export pt format model to aot format ,like AOTPackageLoader , now only could aot loader , like torch.jit.trace and torch.jit.script also important

[mullerhai]

package torch;

import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.global.torch;
import org.bytedeco.javacpp.PointerScope;
import java.util.*;
import java.io.*;

/**
 * JitFullPipelineTest - ?? JIT Pipeline ??
 *
 * ??: SequentialImpl ?? ? save(archive) ? save_to(pt) ? torch.load ? freeze ? optimize
 *
 * ?? API:
 * - Module.save(OutputArchive)      ? ?????
 * - OutputArchive.save_to(filename) ? ?? .pt ??
 * - torch.load(filename)            ? ??? JitModule (torch.jit.load)
 * - torch.freeze(JitModule)         ? ?????
 * - torch.optimize_for_inference    ? AOT ????
 */
public class JitFullPipelineTest {

    static class Result {
        String name; boolean passed; long durationMs; String detail;
        Result(String n, boolean p, long d, String det) {
            name = n; passed = p; durationMs = d; detail = det;
        }
    }

    public static void main(String[] args) throws Exception {
        System.out.println("?????????????????????????????????????????????????????????????????");
        System.out.println("?    JavaCPP-PyTorch JIT Full Pipeline Test                ?");
        System.out.println("?????????????????????????????????????????????????????????????????");
        System.out.println();

        List<Result> results = new ArrayList<>();

        // === ?? Pipeline ?? ===

        // 1. Module.save + save_to ????
        results.add(testModuleSave());

        // 2. InputArchive.load_from + Module.load ????
        results.add(testModuleLoad());

        // 3. ?? Pipeline: save ? load ? freeze ? optimize
        results.add(testFullPipeline());

        // 4. ????????
        results.add(testTrainSaveLoad());

        // === ?????? ===

        // 5. ??????
        results.add(testInferencePerf());

        // 6. ????
        results.add(testModelClone());

        // 7. ????
        results.add(testBatchInference());

        // 8. ????
        results.add(testTraining());

        // ????
        System.out.println("\n???????????????????????????????????????????????????????????????");
        System.out.println("BENCHMARK RESULTS");
        System.out.println("???????????????????????????????????????????????????????????????");

        int passed = 0, failed = 0;
        for (Result r : results) {
            String s = r.passed ? "? PASS" : "? FAIL";
            System.out.printf("%-45s %s (%d ms)%n", r.name, s, r.durationMs);
            if (r.detail != null && !r.detail.isEmpty()) {
                System.out.println("  Detail: " + r.detail);
            }
            if (r.passed) passed++; else failed++;
        }

        System.out.println("\n???????????????????????????????????????????????????????????????");
        System.out.printf("TOTAL: %d/%d PASSED%n", passed, passed + failed);
        System.out.println("???????????????????????????????????????????????????????????????");
        System.exit(failed > 0 ? 1 : 0);
    }

    // 1. Module.save + OutputArchive.save_to
    static Result testModuleSave() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";
        String path = "/tmp/jit_pipeline_model.pt";

        try (PointerScope scope = new PointerScope()) {
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(784, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 10));
            model.eval();

            // ??
            OutputArchive archive = new OutputArchive();
            model.save(archive);
            archive.save_to(path);
            ok = new File(path).exists();
            detail = "saved to " + path + ", exists: " + ok;

            archive.close();
            model.close();
        } catch (Exception e) {
            ok = false;
            detail = "exception: " + e.getMessage();
        }
        return new Result("Module.save + save_to", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // 2. InputArchive.load_from + Module.load
    static Result testModuleLoad() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";
        String path = "/tmp/jit_pipeline_model.pt";

        try (PointerScope scope = new PointerScope()) {
            if (!new File(path).exists()) {
                detail = "model file not found: " + path;
                return new Result("InputArchive.load + Module.load", ok, (System.nanoTime() - start) / 1_000_000, detail);
            }

            SequentialImpl model = new SequentialImpl();
            InputArchive archive = new InputArchive();
            archive.load_from(path);
            model.load(archive);
            ok = !model.isNull();

            if (ok) {
                Tensor in = torch.randn(new long[]{2, 784});
                Tensor out = model.forward(in);
                ok = !out.isNull();
                detail = "loaded, forward valid: " + ok;
                out.close(); in.close();
            }

            archive.close();
            model.close();
        } catch (Exception e) {
            detail = "exception: " + e.getMessage();
        }
        return new Result("InputArchive.load + Module.load", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // 3. ?? Pipeline: save ? load ? freeze ? optimize
    static Result testFullPipeline() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";
        String path = "/tmp/jit_pipeline_full.pt";

        try (PointerScope scope = new PointerScope()) {
            // 1. ???????
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(784, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 10));

            Optimizer opt = new Adam(model.parameters(), new AdamOptions(1e-3));
            for (int i = 0; i < 3; i++) {
                Tensor x = torch.randn(new long[]{4, 784});
                Tensor y = torch.randint(10, new long[]{4});
                Tensor out = model.forward(x);
                Tensor loss = torch.cross_entropy(out, y);
                opt.zero_grad();
                loss.backward();
                opt.step();
                loss.close(); out.close(); y.close(); x.close();
            }
            model.eval();
            opt.close();

            // 2. ??
            OutputArchive archive = new OutputArchive();
            model.save(archive);
            archive.save_to(path);
            archive.close();
            model.close();

            // 3. ? torch.load ??? JitModule
            JitModule loaded = torch.load(path);
            ok = loaded != null && !loaded.isNull();
            detail = "loaded JitModule: " + ok;

            if (!ok) {
                return new Result("Full Pipeline (save?load?freeze?optimize)", ok, (System.nanoTime() - start) / 1_000_000, detail);
            }

            // 4. freeze
            JitModule frozen = torch.freeze(loaded);
            ok = frozen != null && !frozen.isNull();
            detail += ", frozen: " + ok;
            loaded.close();

            if (!ok) {
                return new Result("Full Pipeline", ok, (System.nanoTime() - start) / 1_000_000, detail);
            }

            // 5. optimize_for_inference
            JitModule optimized = torch.optimize_for_inference(frozen);
            ok = optimized != null && !optimized.isNull();
            detail += ", optimized: " + ok;
            frozen.close();

            if (!ok) {
                return new Result("Full Pipeline", ok, (System.nanoTime() - start) / 1_000_000, detail);
            }

            // 6. ??????
            Tensor in = torch.randn(new long[]{2, 784});
            IValueVector inputs = new IValueVector(new IValue(in));
            IValue out = optimized.forward(inputs);
            ok = !out.isNull();
            detail += ", forward valid: " + ok;

            if (!out.isNull()) out.close();
            inputs.close(); in.close();
            optimized.close();

            detail = "Full pipeline: " + detail;
        } catch (Exception e) {
            detail = "exception: " + e.getMessage();
        }
        return new Result("Full Pipeline (save?load?freeze?optimize)", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // 4. ????????
    static Result testTrainSaveLoad() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";
        String path = "/tmp/train_save_load.pt";

        try (PointerScope scope = new PointerScope()) {
            // ??
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(784, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 10));

            Optimizer opt = new Adam(model.parameters(), new AdamOptions(1e-3));
            for (int i = 0; i < 5; i++) {
                Tensor x = torch.randn(new long[]{8, 784});
                Tensor y = torch.randint(10, new long[]{8});
                Tensor out = model.forward(x);
                Tensor loss = torch.cross_entropy(out, y);
                opt.zero_grad();
                loss.backward();
                opt.step();
                loss.close(); out.close(); y.close(); x.close();
            }
            model.eval();

            // ??
            OutputArchive archive = new OutputArchive();
            model.save(archive);
            archive.save_to(path);
            archive.close();
            opt.close();
            model.close();

            // ????
            SequentialImpl loaded = new SequentialImpl();
            InputArchive iarch = new InputArchive();
            iarch.load_from(path);
            loaded.load(iarch);
            ok = !loaded.isNull();

            // ????
            if (ok) {
                Tensor in = torch.randn(new long[]{4, 784});
                Tensor out = loaded.forward(in);
                ok = !out.isNull();
                detail = "trained?saved?loaded, forward: " + ok;
                out.close(); in.close();
            }

            iarch.close();
            loaded.close();
        } catch (Exception e) {
            detail = "exception: " + e.getMessage();
        }
        return new Result("Train ? Save ? Load", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // 5. ????
    static Result testInferencePerf() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(784, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 128));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(128, 10));
            model.eval();

            for (int i = 0; i < 5; i++) {
                Tensor in = torch.randn(new long[]{32, 784});
                Tensor out = model.forward(in);
                out.close(); in.close();
            }

            long total = 0;
            for (int i = 0; i < 100; i++) {
                Tensor in = torch.randn(new long[]{32, 784});
                long t = System.nanoTime();
                Tensor out = model.forward(in);
                total += (System.nanoTime() - t);
                out.close(); in.close();
            }

            double avgMs = total / 1_000_000.0 / 100;
            ok = true;
            detail = String.format("avg: %.3f ms/iter", avgMs);
            model.close();
        } catch (Exception e) {
            detail = "exception: " + e.getMessage();
        }
        return new Result("inference_performance", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // 6. ????
    static Result testModelClone() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(512, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 128));
            model.eval();

            org.bytedeco.pytorch.Module cloned = model.clone(null);
            ok = cloned != null && !cloned.isNull();

            long cnt = 0;
            var params = cloned.parameters();
            var it = params.begin();
            var end = params.end();
            while (!it.equals(end)) { cnt++; it.increment(); }
            detail = "params: " + cnt;

            cloned.close();
            model.close();
        } catch (Exception e) {
            detail = "exception: " + e.getMessage();
        }
        return new Result("model_clone", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // 7. ????
    static Result testBatchInference() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(784, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 10));
            model.eval();

            long total = 0;
            for (int i = 0; i < 100; i++) {
                Tensor in = torch.randn(new long[]{1, 784});
                long t = System.nanoTime();
                Tensor out = model.forward(in);
                total += (System.nanoTime() - t);
                out.close(); in.close();
            }

            double avgMs = total / 1_000_000.0 / 100;
            ok = true;
            detail = String.format("avg: %.3f ms/batch", avgMs);
            model.close();
        } catch (Exception e) {
            detail = "exception: " + e.getMessage();
        }
        return new Result("batch_inference", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }

    // 8. ????
    static Result testTraining() {
        long start = System.nanoTime();
        boolean ok = false;
        String detail = "";

        try (PointerScope scope = new PointerScope()) {
            SequentialImpl model = new SequentialImpl();
            model.push_back(new LinearImpl(784, 256));
            model.push_back(new ReLUImpl());
            model.push_back(new LinearImpl(256, 10));

            Optimizer opt = new Adam(model.parameters(), new AdamOptions(1e-3));
            for (int i = 0; i < 5; i++) {
                Tensor x = torch.randn(new long[]{16, 784});
                Tensor y = torch.randint(10, new long[]{16});
                Tensor out = model.forward(x);
                Tensor loss = torch.cross_entropy(out, y);
                opt.zero_grad();
                loss.backward();
                opt.step();
                loss.close(); out.close(); y.close(); x.close();
            }

            ok = true;
            detail = "trained 5 steps";
            opt.close();
            model.close();
        } catch (Exception e) {
            detail = "exception: " + e.getMessage();
        }
        return new Result("training", ok, (System.nanoTime() - start) / 1_000_000, detail);
    }
}

now we only could do like that

[HGuillemet]

Looks like we'd need to add #include <torch/csrc/jit/serialization/export.h> to torch_include.h.

@HGuillemet Any reason in particular why it's not there already?

The list of include files to parse was generated by walking though the tree of includes, starting from root includes (torch.h and some others: see include_list.pl) , and this file were never reached. It can be added to the root includes if it's part of the C++ API.

[mullerhai]

Looks like we'd need to add #include <torch/csrc/jit/serialization/export.h> to torch_include.h.
@HGuillemet Any reason in particular why it's not there already?

The list of include files to parse was generated by walking though the tree of includes, starting from root includes (torch.h and some others: see include_list.pl) , and this file were never reached. It can be added to the root includes if it's part of the C++ API.

oh, let's try to add it , I think it could make a test on the next snapshot @saudet