Omega-AI: a deep learning framework written in Java from scratch

Written in Java because that is what the author knew

The origin story is in the README's preface, and it explains the project's character better than any feature list would.

The author says that since 2016 they have studied deep learning in their spare time, and because of their work the language they are most familiar with is Java, so the framework was naturally written in Java. The stated original motivation was to understand more deeply the principles, models and implementations of the algorithms, and to make it easier for Java developers to get into the AI field.

That is a pedagogical framing, and it is the one to hold onto when evaluating the project. Omega-AI is not trying to be a drop-in replacement for an existing Java ML library with a benchmark table. It is a from-scratch implementation whose value is partly pedagogical, and partly that it lets a Java team keep everything in one language and one build.

The project lives under the Dromara organisation, with the repository on GitHub and mirrors on Gitee and GitCode, plus an official site at oaii.cn. The licence is Apache-2.0, the last push to the repository was on 2026-10-17, and there are no GitHub releases, so the version you take is a commit rather than a tag.

The topics attached to the repository mention ai, deeplearning, diffusion, llm, neural-network and yolo, which is a fair summary of the range it covers.

The engine is the product, not the model gallery

The framework is split into an engine and a set of model implementations, and the engine is where the engineering lives.

According to the README, the engine supports automatic differentiation, multithreading, and GPU acceleration with CUDA and cuDNN, and the current engine version supports both CUDA and cuDNN as acceleration paths. It also supports multi-GPU training.

Automatic differentiation is what turns this from a maths library into something you can train with. Hand written backpropagation for a transformer is a project in itself, and having the graph and the gradients handled by the engine is what lets the same code style produce GPT, LLaMA and diffusion models.

The dependency policy is stated unusually precisely. The models and algorithms implemented in the engine use no API and no third party dependency other than the CUDA and cuDNN related packages.

That is a strong claim and it is the one to test first. A framework that pulls in nothing beyond the GPU bindings is easier to audit, easier to port to a new platform and much easier to reason about when something behaves unexpectedly, because there is no unknown layer between your model and the hardware. Whether the implementation is complete enough for a given architecture is a separate question, and the way to answer it is to read the model source.

The repository layout supports the split, with a src directory for the Java code, jni and native directories for the native side of the GPU bindings, and a pom.xml for the Maven build.

The jcuda version has to match the CUDA on your machine

The dependency section is short and it is the most likely source of a first day problem.

Because omega-engine-v4-gpu includes jcuda support, that engine needs a CUDA installation matching the jcuda version. The README gives the worked example: if your machine has CUDA version 11.7.x, then the corresponding omega-engine build needs jcuda 11.7.0.

So the GPU build is not a single artifact. There is a family of them, and the one you depend on has to agree with the CUDA toolkit you installed. The README's quick start makes this the first three steps, and the first of them is to check what you actually have.

nvcc --version

Then install CUDA and cuDNN, with the toolkit archive on the NVIDIA developer site linked from the README. Then bring in, or download, the engine package that corresponds to your CUDA version, with a list of Windows CUDA builds given in the version dependency section of the README.

<dependency>
    <groupId>io.gitee.iangellove</groupId>
    <artifactId>omega-engine-v4-gpu</artifactId>
    <version>win-cu11.7-v1.0-beta</version>
</dependency>

Two things in that snippet deserve attention. The version string is a beta build with the CUDA version encoded in it, win-cu11.7-v1.0-beta, which tells you both the platform and the CUDA series, and that the artifacts are still labelled beta. The groupId, io.gitee.iangellove, is the author's own namespace rather than an organisation namespace, which is worth remembering if you are thinking about the long term provenance of a dependency in a commercial product.

Initialise the GPU context and free the memory when you are done

The example program is short enough to read in one go, and it shows the shape of the API.

public static void main(String[] args) {
    try {
        CUDAModules.initContext();
        CNNTest cnn = new CNNTest();
        cnn.cnnNetwork_cifar10();
    } finally {
        CUDAMemoryManager.free();
    }
}

There are three things to notice. The GPU environment is initialised explicitly with a call that returns a Context object, rather than happening implicitly on first use, which means the cost of bringing up the CUDA context is something you control rather than something that lands in the middle of a training run.

Second, the model is a plain Java object. The example constructs a CNNTest instance and calls a method that runs the CIFAR-10 network. There is no separate training script, no configuration file and no data pipeline layer described here, which is a simplification you should be aware of.

Third, and most importantly for anyone running this repeatedly, the memory is released in a finally block with CUDAMemoryManager.free(). In a JVM, native memory is not collected the way the heap is, so an explicit free is what keeps a long running process or a test suite from walking into an out of memory error on the device. Putting it in a finally block rather than at the end of the method is the difference between freeing the memory when the model finishes and freeing it when the model finishes without an exception.

This shape, initialise, run, free, is the pattern to follow in your own code, and it scales to the multi-GPU case because the engine supports training on several devices.

JVM heap flags are part of the model definition

There is a short section on system parameters that is easy to skim past and expensive to ignore.

The README explains that because the parameters for training a vgg16 model are relatively large, the JVM memory has to be adjusted when deploying the project, and it gives an example.

-Xmx20480m -Xms20480m -Xmn10240m

That is a 20 GB maximum heap, a 20 GB initial heap and a 10 GB young generation, which tells you two things at once. First, this is not a framework you will run inside a small container. Second, the memory requirement is a function of the model, not of the framework, so the flags are effectively part of the model definition.

The second point generalises usefully. A framework that trains a large model in Java has to hold the model, the activations and the optimiser state on the managed heap, and the size of the young generation matters because of how much garbage the training loop produces. If you switch models, or switch from training to inference, these numbers change.

The practical consequence is that a deployment checklist for a model built on this engine should include the heap settings alongside the model artefact itself, and that a team moving between models should expect to retune them rather than assume one set of flags works for everything.

A model list that runs from YOLO to diffusion

The framework supports BP neural networks, convolutional neural networks, recurrent neural networks, vgg16, resnet, yolo, lstm, transformer, gpt, llama, diffusion and stable diffusion.

The breadth is the point for a framework that wants to be a single dependency for a team, but breadth is also where the risk lies. Having a class called yolo in the tree tells you the project has an implementation of a YOLO variant. It does not tell you which variant, which input sizes, which augmentation, whether the weights are included or whether you bring your own, or how close the numbers come to a reference implementation.

The demos give a partial answer to those questions, and the demo list is a good guide to what is actually runnable. The convolutional series includes a handwritten digit recognition demo on MNIST, served from what looks like a small web endpoint. The YOLO series includes a car tracking demo with DeepSort, a mask wearing detection demo on yolov3, a safety helmet wearing detection demo on yolov3, and a yolov7 smart refrigerator product recognition demo. The GAN series includes handwritten digit generation and anime avatar generation with DCGAN. The sequence models include an RNN novel generator and a SEQ2SEQ English translator. The GPT series includes a miniature GPT-2 architecture novel generator.

The diffusion work is the most interesting to look at. The Omega-mini-dit demo is a 130M parameter text to image model, described as built on dit_b/1 combined with va-vae, clip, repa, sprint and rms, with the demo images trained on a 2M image and text pair dataset. The README also links a PyTorch version of OmegaDiT, which is the single most useful pointer in the document for anyone who wants to check the implementation, since a reference implementation in the other framework makes correctness arguable rather than a matter of trust.

The demos are the documentation, and the gallery is the test suite

There is no separate documentation site in the repository listing, no wiki and no published API reference. The README itself carries the gallery of demos, and the demos double as the worked examples.

That is a workable arrangement for a project whose audience is a Java developer who wants to see a model run, and it is a weak arrangement for a project whose audience is a team that needs to know what a class does before committing to it.

The honest way to evaluate it is to treat the demo gallery as the specification. Pick the model closest to your problem, look at the code that constructs it, and see whether the pieces you need are there. The repository listing is minimal, with src, jni, native, pom.xml, images and the two READMEs, the English one being README.en.md alongside the Chinese original, so the reading experience is the gallery plus the source.

Two things to keep in mind. The engine build is a beta artifact with the CUDA version embedded in its name, so you are depending on a prerelease package in a production path. And the model implementations depend on nothing outside the CUDA and cuDNN bindings, which is a genuine strength for auditability, but it also means there is no third party layer quietly doing heavy lifting for you, and none of the community tooling that comes with more established frameworks.

The projects use the Apache-2.0 licence, and the source mirrors on Gitee, GitCode and GitHub are all maintained by the Dromara project, which is the sort of arrangement that keeps a project alive if one platform changes its policies.

Editorial conclusion

Omega-AI suits a Java team that wants to build and train models in one language, read the implementation of every layer and algorithm rather than call a black box, and still have GPU acceleration and automatic differentiation, since that is exactly what the engine provides. It does not suit a team that needs a large catalogue of pretrained checkpoints, a mature serving stack, or a framework with years of community documentation behind it, and the dependency situation for the GPU build is something to settle early. Before you commit, check that nvcc --version on your machine matches the CUDA version baked into the jcuda build you depend on, that your JVM heap flags are set for the model you plan to train, since the README's vgg16 example needs a 20 GB heap, and that the engine and model implementations you intend to use are covered by the documentation rather than by the demo gallery alone.

Frequently asked questions

Official sources

1. dromara/Omega-AI on GitHub
2. Issues
3. License: Apache-2.0
4. README