Proposal: Enterprise‑grade, full‑featured PyTorch binding for Java comparable to Python PyTorch ecosystem

[mullerhai]

Hi, @saudet @HGuillemet

Description

This is a proposal to discuss building a more enterprise‑ready, stable, feature‑complete Java PyTorch binding that can closely match the Python PyTorch experience, together with rich upstream/downstream library integrations.

My fork implementation is based on PyTorch 1.5.14 GA release, which resolves many historical bugs and limitations found in existing JavaCPP‑Presets PyTorch bindings. The goal is to bring Java access to a broad ecosystem covering not only core PyTorch, but also more than 50 popular downstream libraries: pandas‑style data utilities, plotting (plot/seaborn), audio, computer vision, NLP, recommendation systems, reinforcement learning modules, and more.

Relevant resources

• Source repository: https://github.com/mullerhai/javacpp-mini-torch
• Maven Central published artifact: https://central.sonatype.com/artifact/io.github.mullerhai/pytorch

Current published dependency set for reference:

<dependency>
    <groupId>io.github.mullerhai</groupId>
    <artifactId>pytorch</artifactId>
    <version>2.13.0-1.5.14-GA-1.10</version>
    <optional>true</optional>
</dependency>
<dependency>
    <groupId>io.github.mullerhai</groupId>
    <artifactId>pytorch</artifactId>
    <version>2.13.0-1.5.14-GA-1.10</version>
    <classifier>linux-x86_64</classifier>
    <optional>true</optional>
</dependency>
<dependency>
    <groupId>io.github.mullerhai</groupId>
    <artifactId>pytorch</artifactId>
    <version>2.13.0-1.5.14-GA-1.10</version>
    <classifier>linux-x86_64-gpu</classifier>
    <optional>true</optional>
</dependency>
<dependency>
    <groupId>io.github.mullerhai</groupId>
    <artifactId>torch-gpu-linux</artifactId>
    <version>13.3-9.25-1.5.14-GA-1.10</version>
    <type>pom</type>
</dependency>
<dependency>
    <groupId>io.github.mullerhai</groupId>
    <artifactId>torch-data</artifactId>
    <version>1.5.14-GA-1.10</version>
</dependency>

y objectives & discussion topics

1. Enterprise‑grade stability: Fix long‑standing bugs, memory‑safety issues, native resource lifecycle problems, JVM crash scenarios, improve production‑readiness for Java backend services.
2. Feature parity with Python PyTorch: Cover missing operators, tensor APIs, autograd, JIT / torch.save / torch.load workflows, training‑loop primitives.
3. Ecosystem expansion: Integrate common downstream components:
   • Data processing: torch‑data, pandas‑like data abstractions
   • Domains: audio, vision, NLP, recommendation, reinforcement learning
   • Visualization & analysis: plot / seaborn‑style bindings
4. Distribution model discussion:
   • Multi‑arch native binary packaging (CPU / GPU classifier split, platform‑specific POM BOM as torch‑gpu‑linux)
   • Maven dependency ergonomics, optional native artifacts, BOM import pattern, reduce user configuration friction
5. Collaboration path:
   • Whether parts of this work can be upstreamed into JavaCPP‑Presets, or maintained as a separate companion project under the ByteDeco umbrella.
   • Compatibility strategy: maintain interoperability with existing JavaCPP types vs potential new API surface.
   • CI/CD, test coverage, release cadence for enterprise usage scenarios.

some code

/*
 * Java port of /home/muller/下载/muTorch-main/src/main/scala/TrainerV2.scala
 *
 *   Trainer V2 - MicroLens / Transformer_DCN with the new PyTorch DataFrame stack.
 *
 *   Differences from the V1 Trainer:
 *     - DataFrame.readParquet for train / valid / item_info
 *     - DataFrameDataset.builder(...)  ->  features(...).sequenceFeature(...).labels(...).build()
 *     - DataFrameDataLoader (pure-Java) for named multi-feature batches
 *     - Tqdm progress bars instead of hand-rolled log lines
 *     - Per-batch PointerScope so JavaCPP native Tensor wrappers are released
 *       immediately (without this, RSS grows and training gets progressively slower)
 *
 *   Defaults: batch=256, epochs=10, lr=5e-4 (safe on CUDA).
 *
 *   Args:
 *     0  dataRoot        (default /home/muller/下载/pytorch-2025MMCTR/WWW2025_MMCTR/data/MicroLens_1M_x1)
 *     1  batchSize       (default 256)
 *     2  numEpochs       (default 10)
 *     3  learningRate    (default 5e-4)
 *     4  maxTrainBatches (default 0 = all)
 *     5  maxValidBatches (default 0 = all)
 *     6  logEvery        (default 10)
 *     7  trainSampleFrac (default 1.0)
 *     8  profileFirstN   (default 5)
 *     9  device          (default "cuda"; falls back to "mps" / "cpu")
 *
 *   Run:
 *     java -cp <classpath> samples.recommend.TrainerV2Java [args]
 */
package example;

import org.bytedeco.javacpp.Loader;
import org.bytedeco.javacpp.PointerScope;
import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.dataframe.Column;
import org.bytedeco.pytorch.dataframe.DataFrame;
import org.bytedeco.pytorch.dataframe.dataset.DataFrameDataLoader;
import org.bytedeco.pytorch.dataframe.dataset.DataFrameDataset;
import org.bytedeco.pytorch.global.torch;
import org.bytedeco.pytorch.global.torch.ScalarType;
import org.bytedeco.pytorch.optim.Adam;
import org.bytedeco.pytorch.optim.options.AdamOptions;
import org.bytedeco.pytorch.plot.tqdm.ProgressBarColor;
import org.bytedeco.pytorch.plot.tqdm.Tqdm;
import org.bytedeco.pytorch.plot.tqdm.TqdmBar;
import org.bytedeco.pytorch.recommend.DeviceSupport;
import org.bytedeco.pytorch.recommend.basic.metrics.RankingMetrics;
import org.bytedeco.pytorch.recommend.models.ranking.TransformerDCN;

import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

public final class TrainerV2Java {

    private static final int    PRETRAIN_DIM = 128;
    private static final int    TAGS_LEN     = 5;
    private static final int    MAX_SEQ_LEN  = 64;

    private TrainerV2Java() {}

    // ------------------------------------------------------------------
    //  Batch: same shape as Scala torchrec.train.Batch
    // ------------------------------------------------------------------

    public static final class Batch {
        public final Tensor itemSeq;
        public final Tensor itemIds;
        public final Tensor likesLevel;
        public final Tensor viewsLevel;
        public final Tensor mask;
        public final Tensor label;
        public final int    batchSize;

        public Batch(Tensor itemSeq, Tensor itemIds, Tensor likesLevel,
                     Tensor viewsLevel, Tensor mask, Tensor label, int batchSize) {
            this.itemSeq    = itemSeq;
            this.itemIds    = itemIds;
            this.likesLevel = likesLevel;
            this.viewsLevel = viewsLevel;
            this.mask       = mask;
            this.label      = label;
            this.batchSize  = batchSize;
        }

        /** Move all tensors to {@code device}, same as Scala {@code Batch.to(device)}. */
        public Batch to(Device device) {
            TensorOptions opts = new TensorOptions(device);
            MemoryFormatOptional mf = new MemoryFormatOptional();
            return new Batch(
                    itemSeq.to(opts, false, false, mf),
                    itemIds.to(opts, false, false, mf),
                    likesLevel.to(opts, false, false, mf),
                    viewsLevel.to(opts, false, false, mf),
                    mask.to(opts, false, false, mf),
                    label.to(opts, false, false, mf),
                    batchSize);
        }
    }

    // ------------------------------------------------------------------
    //  Ranking metrics - same logic as Scala RankingMetrics
    // ------------------------------------------------------------------

    public static final class Metrics {

        private Metrics() {}

        /** Wall-clock ms since {@code t0Nano} (from {@link System#nanoTime()}). */
        public static double msSince(long t0Nano) {
            return (System.nanoTime() - t0Nano) / 1e6;
        }

        /** Postfix map for {@code TqdmBar.set_postfix(...)}. */
        public static Map<String, Object> postfix(double loss, Double aucOpt, double sps) {
            Map<String, Object> m = new LinkedHashMap<>();
            m.put("loss", String.format("%.4f", loss));
            if (aucOpt != null) m.put("auc", String.format("%.3f", aucOpt));
            m.put("sps", String.format("%,d", (long) sps));
            return m;
        }
    }

    // ------------------------------------------------------------------
    //  Device resolution
    // ------------------------------------------------------------------

    /**
     * Resolves which device to use. Order matches the original Scala, but
     * is also overridable via the 9th CLI argument ("cuda" / "mps" / "cpu" / "auto").
     */
    private static Device resolveDevice(String requested) {
        String req = requested == null ? "cuda" : requested.toLowerCase();
        // Force DeviceSupport backend resolution before any tensor allocation.
        switch (req) {
            case "cuda":
                if (DeviceSupport.cudaAvailable()) {
                    DeviceSupport.setDevice("cuda");
                    System.out.println("  Device       : CUDA");
                    return new Device(torch.kCUDA());
                }
                System.out.println("  Device       : CUDA requested but unavailable; falling back to auto");
                break;
            case "mps":
                if (torch.hasMPS()) {
                    DeviceSupport.setDevice("mps");
                    System.out.println("  Device       : MPS");
                    return new Device(torch.kMPS());
                }
                System.out.println("  Device       : MPS requested but unavailable; falling back to auto");
                break;
            case "cpu":
                DeviceSupport.setDevice("cpu");
                System.out.println("  Device       : CPU (explicit)");
                return new Device(torch.kCPU());
            case "auto":
            default:
                break;
        }
        String backend = DeviceSupport.backend();
        Device dev = new Device("cpu".equals(backend) ? torch.kCPU() :
                "cuda".equals(backend) ? torch.kCUDA() : torch.kMPS());
        System.out.println("  Device       : " + backend.toUpperCase());
        return dev;
    }

    // ------------------------------------------------------------------
    //  item_info tables
    // ------------------------------------------------------------------

    /** Flattened item tables keyed by row index == item_id (row 0 = padding). */
    public static final class ItemTables {
        public final int      numItems;
        public final float[]  pretrainedEmbFlat;  // [numItems * 128]
        public final long[]   itemTagsFlat;       // [numItems * 5]

        public ItemTables(int numItems, float[] pretrainedEmbFlat, long[] itemTagsFlat) {
            this.numItems          = numItems;
            this.pretrainedEmbFlat = pretrainedEmbFlat;
            this.itemTagsFlat      = itemTagsFlat;
        }
    }

    /** Symmetric to Scala {@code loadItemTables}. */
    private static ItemTables loadItemTables(String itemPath) throws Exception {
        System.out.print("Loading item_info via DataFrame: " + itemPath + " ... ");
        System.out.flush();
        DataFrame df = DataFrame.readParquet(itemPath);
        int n = df.rowCount();
        System.out.printf("%,d items%n", n);

        float[] embFlat  = new float[n * PRETRAIN_DIM];
        long[]  tagsFlat = new long[n * TAGS_LEN];

        Column embCol  = df.hasColumn("item_emb_d128") ? df.column("item_emb_d128") : null;
        Column tagsCol = df.hasColumn("item_tags")     ? df.column("item_tags")     : null;

        for (int i = 0; i < n; i++) {
            if (embCol  != null) fillFloatRow(embFlat,  i * PRETRAIN_DIM, PRETRAIN_DIM, embCol.get(i));
            if (tagsCol != null) fillLongRow (tagsFlat, i * TAGS_LEN,     TAGS_LEN,     tagsCol.get(i));
        }
        return new ItemTables(n, embFlat, tagsFlat);
    }

    /** Same logic as Scala {@code fillLongRow}. */
    private static void fillLongRow(long[] dest, int off, int dim, Object v) {
        if (v == null) return;
        if (v instanceof long[]) {
            long[] a = (long[]) v;
            System.arraycopy(a, 0, dest, off, Math.min(dim, a.length));
        } else if (v instanceof int[]) {
            int[] a = (int[]) v;
            int m = Math.min(dim, a.length);
            for (int k = 0; k < m; k++) dest[off + k] = a[k];
        } else if (v instanceof double[]) {
            double[] a = (double[]) v;
            int m = Math.min(dim, a.length);
            for (int k = 0; k < m; k++) dest[off + k] = (long) a[k];
        } else if (v instanceof float[]) {
            float[] a = (float[]) v;
            int m = Math.min(dim, a.length);
            for (int k = 0; k < m; k++) dest[off + k] = (long) a[k];
        } else if (v instanceof List) {
            List<?> list = (List<?>) v;
            int m = Math.min(dim, list.size());
            for (int k = 0; k < m; k++) {
                Object e = list.get(k);
                if (e instanceof Number) dest[off + k] = ((Number) e).longValue();
            }
        } else if (v instanceof Number) {
            dest[off] = ((Number) v).longValue();
        }
    }

    /** Same logic as Scala {@code fillFloatRow}. */
    private static void fillFloatRow(float[] dest, int off, int dim, Object v) {
        if (v == null) return;
        if (v instanceof float[]) {
            float[] a = (float[]) v;
            System.arraycopy(a, 0, dest, off, Math.min(dim, a.length));
        } else if (v instanceof double[]) {
            double[] a = (double[]) v;
            int m = Math.min(dim, a.length);
            for (int k = 0; k < m; k++) dest[off + k] = (float) a[k];
        } else if (v instanceof long[]) {
            long[] a = (long[]) v;
            int m = Math.min(dim, a.length);
            for (int k = 0; k < m; k++) dest[off + k] = (float) a[k];
        } else if (v instanceof int[]) {
            int[] a = (int[]) v;
            int m = Math.min(dim, a.length);
            for (int k = 0; k < m; k++) dest[off + k] = (float) a[k];
        } else if (v instanceof List) {
            List<?> list = (List<?>) v;
            int m = Math.min(dim, list.size());
            for (int k = 0; k < m; k++) {
                Object e = list.get(k);
                if (e instanceof Number) dest[off + k] = ((Number) e).floatValue();
            }
        } else if (v instanceof Number) {
            dest[off] = ((Number) v).floatValue();
        }
    }

    // ------------------------------------------------------------------
    //  DataFrame -> Dataset -> DataLoader
    // ------------------------------------------------------------------

    private static DataFrameDataset buildDataset(String path, String name,
                                                 double sampleFrac, long sampleSeed) throws Exception {
        System.out.print("Loading " + name + " via DataFrame: " + path + " ... ");
        System.out.flush();
        DataFrame raw = DataFrame.readParquet(path);
        int fullN = raw.rowCount();
        DataFrame df;
        if (sampleFrac < 1.0 && sampleFrac > 0.0) {
            DataFrame sampled = raw.sampleFrac(sampleFrac, sampleSeed);
            System.out.printf("%,d rows -> sampleFrac=%.4f -> %,d rows%n",
                    fullN, sampleFrac, sampled.rowCount());
            df = sampled;
        } else {
            System.out.printf("%,d rows (full)%n", fullN);
            df = raw;
        }
        df.printSchema();
        df.show(3);

        // labelsAsLong(false) -> float labels for BCEWithLogitsLoss
        return DataFrameDataset.builder(df)
                .features("user_id", "item_id", "likes_level", "views_level")
                .sequenceFeature("item_seq")
                .labels("label")
                .labelsAsLong(false)
                .build();
    }

    private static DataFrameDataLoader buildLoader(DataFrameDataset ds, int batchSize,
                                                   boolean shuffle, long seed) {
        return ds.dataloader()
                .batchSize(batchSize)
                .shuffle(shuffle)
                .dropLast(false)
                .seed(seed)
                .build();
    }

    // ------------------------------------------------------------------
    //  DataFrameDataLoader.Batch -> our Batch
    // ------------------------------------------------------------------

    /**
     * Map named DataFrame features onto the model {@link Batch}:
     *   item_seq          -> long [B,S]
     *   item_id           -> long [B]
     *   likes_level       -> long [B]
     *   views_level       -> long [B]
     *   mask = (item_seq > 0).float()
     *   label             -> float [B]
     */
    private static Batch toModelBatch(DataFrameDataLoader.Batch b, Device device) {
        Tensor itemSeq    = b.feature("item_seq").to(ScalarType.Long).contiguous();
        Tensor itemIds    = b.feature("item_id").to(ScalarType.Long).contiguous();
        Tensor likesLevel = b.feature("likes_level").to(ScalarType.Long).contiguous();
        Tensor viewsLevel = b.feature("views_level").to(ScalarType.Long).contiguous();
        Tensor mask       = itemSeq.gt(new Scalar(0L)).to(ScalarType.Float).contiguous();

        Tensor labelRaw = b.labels();
        Tensor label;
        if (labelRaw == null) {
            label = torch.zeros(new long[]{b.size()}, new TensorOptions(ScalarType.Float));
        } else {
            label = labelRaw.to(ScalarType.Float).reshape(new long[]{b.size()}).contiguous();
        }
        Batch batch = new Batch(itemSeq, itemIds, likesLevel, viewsLevel,
                                mask, label, b.size());
        return batch.to(device);
    }

    // ------------------------------------------------------------------
    //  Model budget print
    // ------------------------------------------------------------------

    private static void reportModelBudget(TransformerDCN model, int batchSize) {
        TensorVector params = model.parameters();
        long n = params.size();
        long nParams = 0L;
        for (long i = 0; i < n; i++) {
            nParams += params.get(i).__dispatch_numel();
        }
        double weightMb = nParams * 4.0 / (1024.0 * 1024.0);
        double adamMb   = weightMb * 3.0;
        long d = model.dcnInDim();
        long crossParams = 3L * d * d + 3L * d;
        double crossMb = crossParams * 4.0 / (1024.0 * 1024.0);
        double actMb   = (double) batchSize * (double) d * 4.0 * 8.0 / (1024.0 * 1024.0);
        System.out.printf("  Params         : %,d  (~%.0f MB fp32 weights)%n",      nParams,    weightMb);
        System.out.printf("  Adam state est : ~%.0f MB (weights + m + v)%n",        adamMb);
        System.out.printf("  CrossNet alone : d=%d  ~%,d params (~%.0f MB)  act@bs=%d ~%.0f MB%n",
                d, crossParams, crossMb, batchSize, actMb);
    }

    /**
     * Warm up Adam state tensors OUTSIDE any PointerScope, so {@code step()}
     * allocates exp_avg / exp_avg_sq lazily before per-batch scopes start.
     */
    private static void warmAdamState(TransformerDCN model, Adam optimizer,
                                      Device device, int batchSize) {
        System.out.print("Warming Adam state (outside PointerScope) ... ");
        System.out.flush();
        long B = Math.min(batchSize, 4);
        long S = MAX_SEQ_LEN;
        Tensor itemSeq = torch.zeros(new long[]{B, S}, new TensorOptions(ScalarType.Long));
        Tensor itemIds = torch.zeros(new long[]{B},    new TensorOptions(ScalarType.Long));
        Tensor likes   = torch.zeros(new long[]{B},    new TensorOptions(ScalarType.Long));
        Tensor views   = torch.zeros(new long[]{B},    new TensorOptions(ScalarType.Long));
        Tensor mask    = torch.ones (new long[]{B, S}, new TensorOptions(ScalarType.Float));
        Tensor label   = torch.zeros(new long[]{B},    new TensorOptions(ScalarType.Float));
        Batch batch = new Batch(itemSeq, itemIds, likes, views, mask, label, (int) B).to(device);

        model.train(true);
        optimizer.zero_grad();
        Tensor out  = model.forward(batch.itemSeq, batch.itemIds, batch.mask, batch.likesLevel, batch.viewsLevel);
        Tensor loss = model.computeLoss(batch.label, out);
        loss.backward();
        optimizer.step();
        optimizer.zero_grad();
        System.out.println("done.");
    }

    // ------------------------------------------------------------------
    //  main
    // ------------------------------------------------------------------

    public static void main(String[] args) throws Exception {
        String dataRoot = args.length > 0 ? args[0]
                : "/home/muller/下载/极光doc/doc/code/cpp/VideoMMCTR/data/MicroLens_1M_x1";
        int    batchSize       = args.length > 1 ? Integer.parseInt(args[1]) : 256;
        int    numEpochs       = args.length > 2 ? Integer.parseInt(args[2]) : 10;
        double learningRate    = args.length > 3 ? Double.parseDouble(args[3]) : 5e-4;
        int    maxTrainBatches = args.length > 4 ? Integer.parseInt(args[4]) : 0;
        int    maxValidBatches = args.length > 5 ? Integer.parseInt(args[5]) : 0;
        int    logEvery        = args.length > 6 ? Math.max(1, Integer.parseInt(args[6])) : 10;
        double trainSampleFrac = args.length > 7 ? Double.parseDouble(args[7]) : 1.0;
        int    profileFirstN   = args.length > 8 ? Integer.parseInt(args[8]) : 5;
//        String requestedDevice = args.length > 9 ? args[9] : "cuda";
        String requestedDevice =  "cuda";

        System.out.println("============================================================");
        System.out.println(" Transformer_DCN - MicroLens Training V2 (DataFrame + tqdm) [JAVA]");
        System.out.println("============================================================");
        System.out.println("  Data root    : " + dataRoot);
        System.out.println("  Batch size   : " + batchSize);
        System.out.println("  Epochs       : " + numEpochs);
        System.out.println("  Learning rate: " + learningRate);
        System.out.println("  Max train batches (0=all): " + maxTrainBatches);
        System.out.println("  Max valid batches (0=all): " + maxValidBatches);
        System.out.println("  Log every    : " + logEvery + " batch(es)  (batch-AUC only on log steps)");
        System.out.println("  Train sample : " + (trainSampleFrac * 100.0) + "% of full train set");
        System.out.println("  Profile first: " + profileFirstN + " train batches (ms breakdown)");
        System.out.println("  Device       : " + requestedDevice);

        Loader.load(torch.class);
        torch.manual_seed(20242025L);
        Device device = resolveDevice(requestedDevice);
        System.out.println();

        String trainPath = dataRoot + "/train.parquet";
        String validPath = dataRoot + "/valid.parquet";
        String itemPath  = dataRoot + "/item_info.parquet";

        // 1) item tables (still needed by FeatureEmbedding)
        ItemTables itemTables = loadItemTables(itemPath);

        // 2) DataFrame -> Dataset -> DataLoader
        DataFrameDataset trainDs = buildDataset(trainPath, "train", trainSampleFrac, 20242025L);
        DataFrameDataset validDs = buildDataset(validPath, "valid", 1.0,            0L);

        DataFrameDataLoader trainLoader = buildLoader(trainDs, batchSize, true,  20242025L);
        DataFrameDataLoader validLoader = buildLoader(validDs, batchSize, false, 0L);

        int trainTotalBatches = maxTrainBatches > 0
                ? Math.min(maxTrainBatches, trainLoader.numBatches())
                : trainLoader.numBatches();
        int validTotalBatches = maxValidBatches > 0
                ? Math.min(maxValidBatches, validLoader.numBatches())
                : validLoader.numBatches();

        System.out.printf("  Train samples     : %,d%n", trainDs.size());
        System.out.printf("  Valid samples     : %,d%n", validDs.size());
        System.out.printf("  Train steps/epoch : %,d  (~%,d samples)%n",
                trainTotalBatches, (long) trainTotalBatches * batchSize);
        System.out.printf("  Valid steps       : %,d%n", validTotalBatches);
        System.out.println();

        System.out.print("Building model ... ");
        System.out.flush();
        long nItems = Math.max(itemTables.numItems, 91718L);
        TransformerDCN model = new TransformerDCN(
                /* itemVocabSize     */ nItems,
                /* embDim            */ 64L,
                /* pretrainDim       */ PRETRAIN_DIM,
                /* embDimPretrain    */ 128L,
                /* likesVocabSize    */ 11L,
                /* viewsVocabSize    */ 11L,
                /* tagsVocabSize     */ 11740L,
                /* numItems          */ nItems,
                /* tagsLen           */ TAGS_LEN,
                /* pretrainedEmbFlat */ itemTables.pretrainedEmbFlat,
                /* itemTagsFlat      */ itemTables.itemTagsFlat,
                /* numHeads          */ 1L,
                /* transformerLayers */ 2,
                /* transformerDropout*/ 0.2,
                /* dimFeedforward    */ 256L,
                /* firstKCols        */ 16,
                /* concatMaxPool     */ true,
                /* dcnCrossLayers    */ 3,
                /* dcnHiddenUnits    */ new long[]{1024L, 512L, 256L},
                /* mlpHiddenUnits    */ new long[]{64L, 32L},
                /* netDropout        */ 0.2,
                /* device            */ device.type().toString().toLowerCase());
        model.to(device, false);
        model.train(true);
        System.out.println("done.");

        System.out.println("  Expected dcnInDim ~ 9088");
        reportModelBudget(model, batchSize);
        System.out.println();

        // Build optimizer OUTSIDE any PointerScope so param list / state live for the run.
        TensorVector params = model.parameters();
        Adam optimizer = new Adam(params, new AdamOptions(learningRate));
        warmAdamState(model, optimizer, device, batchSize);
        System.out.println();

        // ---- epoch loop ----
        for (int epoch = 0; epoch < numEpochs; epoch++) {

            // ---------- TRAIN ----------
            model.train(true);
            long trainExpected = maxTrainBatches > 0
                    ? (long) maxTrainBatches * batchSize
                    : trainDs.size();
            RankingMetrics.Accumulator trainAcc = new RankingMetrics.Accumulator(
                    (int) Math.min(Integer.MAX_VALUE, Math.max(4096L, trainExpected)));
            int  batchCount  = 0;
            long samplesDone = 0L;
            long t0          = System.nanoTime();

            Tqdm.write(String.format("--- Epoch %d/%d TRAIN  batch_size=%d  steps=%,d ---",
                    epoch + 1, numEpochs, batchSize, trainTotalBatches));

            Iterator<DataFrameDataLoader.Batch> trainIt = trainLoader.iterator();
            try (TqdmBar<DataFrameDataLoader.Batch> trainBar = Tqdm.of(trainIt, trainTotalBatches)
                    .setDescription("E" + (epoch + 1) + " train")
                    .setUnit("batch")
                    .setMinInterval(0.2)
                    .setAscii(false)
                    .colour(ProgressBarColor.fromName("green").name())) {

                while (trainBar.hasNext()
                        && (maxTrainBatches <= 0 || batchCount < maxTrainBatches)) {
                    boolean doProfile = batchCount < profileFirstN;
                    long tBatch0 = System.nanoTime();

                    PointerScope scope = new PointerScope();
                    double lossVal  = 0.0;
                    int    actualBs = 0;
                    Double aucOpt   = null;
                    try {
                        long tData = System.nanoTime();
                        DataFrameDataLoader.Batch dfBatch = trainBar.next();   // collate attaches to scope
                        Batch batch = toModelBatch(dfBatch, device);
                        actualBs = batch.batchSize;
                        double msData = Metrics.msSince(tData);

                        long tFwd = System.nanoTime();
                        optimizer.zero_grad();
                        Tensor out  = model.forward(batch.itemSeq, batch.itemIds, batch.mask,
                                batch.likesLevel, batch.viewsLevel);
                        Tensor loss = model.computeLoss(batch.label, out);
                        if (doProfile) loss.item_double();
                        double msFwd = Metrics.msSince(tFwd);

                        long tBwd = System.nanoTime();
                        loss.backward();
                        if (doProfile) loss.item_double();
                        double msBwd = Metrics.msSince(tBwd);

                        long tStep = System.nanoTime();
                        optimizer.step();
                        lossVal = loss.item_double();
                        double msStep = Metrics.msSince(tStep);

                        long tMet = System.nanoTime();
                        trainAcc.update(out, batch.label, lossVal, actualBs);
                        batchCount += 1;
                        samplesDone += actualBs;

                        boolean shouldLog =
                                batchCount == 1 ||
                                batchCount == trainTotalBatches ||
                                batchCount % logEvery == 0;
                        if (shouldLog) {
                            RankingMetrics.BatchMetrics bm = RankingMetrics.fromLogits(out, batch.label);
                            aucOpt = bm.auc;
                        }
                        double msMet = Metrics.msSince(tMet);

                        if (doProfile) {
                            double msTot = Metrics.msSince(tBatch0);
                            Tqdm.write(String.format(
                                    "  [profile b%02d] data=%.0fms  fwd=%.0fms  " +
                                    "bwd=%.0fms  step=%.0fms  met=%.0fms  total=%.0fms  bs=%d",
                                    batchCount, msData, msFwd, msBwd, msStep, msMet, msTot, actualBs));
                        }
                    } finally {
                        scope.close();
                    }

                    double elapsed = (System.nanoTime() - t0) / 1e9;
                    double sps     = elapsed > 0 ? samplesDone / elapsed : 0.0;
                    trainBar.set_postfix(Metrics.postfix(lossVal, aucOpt, sps));
                }
            }

            double[] tr = trainAcc.result();
            double trElapsed = (System.nanoTime() - t0) / 1e9;
            Tqdm.write(String.format(
                    "[E%d] TRAIN done  n=%,d  batches=%,d  logloss=%.5f  auc=%.4f  acc=%.4f  %,d samples/s  %.1fs",
                    epoch + 1, (long) tr[3], batchCount, tr[0], tr[1], tr[2],
                    (long) (tr[3] / trElapsed), trElapsed));
            trainAcc.clear();

            // ---------- VALID ----------
            model.eval();
            long validExpected = maxValidBatches > 0
                    ? (long) maxValidBatches * batchSize
                    : validDs.size();
            RankingMetrics.Accumulator validAcc = new RankingMetrics.Accumulator(
                    (int) Math.min(Integer.MAX_VALUE, Math.max(4096L, validExpected)));
            int  vCount   = 0;
            long vSamples = 0L;
            long tv0      = System.nanoTime();

            Tqdm.write(String.format("--- Epoch %d/%d VALID  steps=%,d ---",
                    epoch + 1, numEpochs, validTotalBatches));

            Iterator<DataFrameDataLoader.Batch> validIt = validLoader.iterator();
            try (TqdmBar<DataFrameDataLoader.Batch> validBar = Tqdm.of(validIt, validTotalBatches)
                    .setDescription("E" + (epoch + 1) + " valid")
                    .setUnit("batch")
                    .setMinInterval(0.2)
                    .setAscii(false)
                    .colour(ProgressBarColor.fromName("cyan").name())) {

                while (validBar.hasNext()
                        && (maxValidBatches <= 0 || vCount < maxValidBatches)) {
                    PointerScope scope = new PointerScope();
                    double lossVal = 0.0;
                    int    bs      = 0;
                    try {
                        DataFrameDataLoader.Batch dfBatch = validBar.next();
                        Batch batch = toModelBatch(dfBatch, device);
                        bs = batch.batchSize;
                        Tensor out  = model.forward(batch.itemSeq, batch.itemIds, batch.mask,
                                batch.likesLevel, batch.viewsLevel);
                        Tensor loss = model.computeLoss(batch.label, out);
                        lossVal = loss.item_double();
                        validAcc.update(out, batch.label, lossVal, bs);
                        vCount  += 1;
                        vSamples += bs;
                    } finally {
                        scope.close();
                    }
                    double elapsed = (System.nanoTime() - tv0) / 1e9;
                    double sps     = elapsed > 0 ? vSamples / elapsed : 0.0;
                    validBar.set_postfix(Metrics.postfix(lossVal, null, sps));
                }
            }

            double[] vr = validAcc.result();
            if (vr[3] > 0) {
                double vElapsed = (System.nanoTime() - tv0) / 1e9;
                Tqdm.write(String.format(
                        "[E%d] VALID done  n=%,d  batches=%,d  logloss=%.5f  auc=%.4f  acc=%.4f  %,d samples/s  %.1fs",
                        epoch + 1, (long) vr[3], vCount, vr[0], vr[1], vr[2],
                        (long) (vr[3] / vElapsed), vElapsed));
            }
            validAcc.clear();
        }

        System.out.println("Training complete (V2 / DataFrame + tqdm + PointerScope) [JAVA / "
                + device.type().toString().toLowerCase() + "].");
    }
}