Skip to content

Commit 000ace7

Browse files
authored
Merge pull request #2666 from annietllnd/neural-graphics
Update Model Gym LP
2 parents a1d1f33 + 695635b commit 000ace7

File tree

2 files changed

+171
-2
lines changed

2 files changed

+171
-2
lines changed

content/learning-paths/mobile-graphics-and-gaming/model-training-gym/4-model-explorer.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,6 @@ model-explorer --extensions=vgf_adapter_model_explorer
4444

4545
Use the file browser to open the `.vgf` model exported earlier in your training workflow.
4646

47-
## Wrapping up
47+
Continue to the next section to learn about bringing your own use-cases to the Model Gym.
48+
4849

49-
Through this Learning Path, you’ve learned what neural graphics is and why it matters for game performance. You’ve stepped through the process of training and evaluating an NSS model using PyTorch and the Model Gym, and seen how to export that model into VGF (.vgf) for real-time deployment. You’ve also explored how to visualize and inspect the model’s structure using Model Explorer. You can now explore the Model Training Gym repository for deeper integration and to keep building your skills.
Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
---
2+
title: Defining your own use cases
3+
weight: 6
4+
5+
### FIXED, DO NOT MODIFY
6+
layout: learningpathall
7+
---
8+
9+
## Beyond NSS: working with your own model
10+
11+
While NSS is a powerful demonstration of neural graphics, Model Gym is designed to support custom models for various use cases. You can add your own model architecture, register it with the toolkit, and run the same training, evaluation, and export workflows you used for NSS.
12+
13+
This section walks you through the process of integrating a custom model into Model Gym, using the Python API you used in previous sections.
14+
15+
### Create a custom model
16+
17+
To add a custom model to the Model Gym, you need to:
18+
19+
1. Create a Python class that inherits from `BaseNGModel`
20+
2. Mark it with the `@register_model()` decorator
21+
3. Implement required methods
22+
4. Accept `params` as a constructor argument
23+
24+
The resulting structure looks like this:
25+
26+
```python
27+
from ng_model_gym.core.model.base_ng_model import BaseNGModel
28+
from ng_model_gym.core.model.model_registry import register_model
29+
from ng_model_gym.core.utils.config_model import ConfigModel
30+
31+
@register_model(name="custom_model", version="1")
32+
class CustomModel(BaseNGModel):
33+
def __init__(self, params: ConfigModel):
34+
super().__init__(params)
35+
# Define your model architecture here
36+
37+
def forward(self, x):
38+
# Implement forward pass
39+
pass
40+
```
41+
42+
The `@register_model()` decorator makes your model discoverable by Model Gym. The `name` and `version` parameters will be used later in your configuration file.
43+
44+
### Register your model
45+
46+
For your model to be available in Model Gym, the file defining it must be imported. This triggers the registration process.
47+
48+
Place your model file in the `src/ng_model_gym/usecases/` directory within the Model Gym installation. To find the installed location, make sure the virtual environment you used with the example notebooks is activated. Then, run the following:
49+
50+
```bash
51+
whereis ng_model_gym
52+
```
53+
54+
This should point to your `nb-env` virtual environment. The Model Gym source code in this case will sit in the following directory:
55+
56+
```output
57+
<path-to>/nb-env/lib/python3.12/site-packages/ng_model_gym/usecases
58+
```
59+
60+
Each use case directory must contain an `__init__.py` file. To understand how the configuration files work, revisit the example notebooks for NSS in previous sections.
61+
62+
Verify registration succeeded by listing registered models:
63+
64+
```python
65+
from ng_model_gym.core.model.model_registry import MODEL_REGISTRY
66+
67+
MODEL_REGISTRY.list_registered()
68+
```
69+
70+
### Update your configuration file
71+
72+
You've already worked with the configuration file in the training and evaluation notebooks. To use your custom model, update the same JSON configuration file to reference the model name and version you used when registering. You can generate a new config template with `ng-model-gym init`, or modify an existing one.
73+
74+
Update the model section:
75+
76+
```json
77+
{
78+
"model": {
79+
"name": "custom_model",
80+
"version": "1"
81+
}
82+
}
83+
```
84+
85+
If you're working in Python, set these values programmatically:
86+
87+
```python
88+
import ng_model_gym as ngmg
89+
from pathlib import Path
90+
91+
config = ngmg.load_config_file(Path("config.json"))
92+
config.model.name = "custom_model"
93+
config.model.version = "1"
94+
```
95+
96+
### Run training with your custom model
97+
98+
Once your model is registered and your config is updated, you can use all the standard Model Gym workflows that were covered in previous sections. For example training:
99+
100+
```bash
101+
trained_model_path = ngmg.do_training(config, training_mode=TrainEvalMode.FP32)
102+
```
103+
104+
Your custom model will go through the same pipeline as NSS: training, quantization, evaluation, and export to `.vgf` format.
105+
106+
### Add custom datasets
107+
108+
Similar to models, you can register custom datasets by marking your dataset class with the `@register_dataset()` decorator:
109+
110+
```python
111+
from torch.utils.data import Dataset
112+
from ng_model_gym.core.data.dataset_registry import register_dataset
113+
114+
@register_dataset(name="custom_dataset", version="1")
115+
class CustomDataset(Dataset):
116+
def __init__(self, config):
117+
# Initialize your dataset
118+
pass
119+
120+
def __len__(self):
121+
# Return dataset size
122+
pass
123+
124+
def __getitem__(self, idx):
125+
# Return a single sample
126+
pass
127+
```
128+
129+
Place the dataset implementation in the same `usecases` directory as your model. Then, update your configuration to use the custom dataset:
130+
131+
```json
132+
{
133+
"dataset": {
134+
"name": "custom_dataset",
135+
"version": "1"
136+
}
137+
}
138+
```
139+
140+
## Summary of self-defined use cases
141+
142+
You can group related models, datasets, and configurations into custom use cases. This is useful when working on a specific neural graphics application. To understand what the final structure should look like, use the NSS use case as a reference implementation. Your use case logic executes when you specify its model and dataset in your configuration file.
143+
144+
## Explore the example notebook
145+
146+
The `neural-graphics-model-gym-examples` repository includes a walkthrough notebook called `custom_model_example.ipynb`. This notebook demonstrates:
147+
148+
- How to structure a custom model class
149+
- Registration and verification steps
150+
- Modifying configuration files
151+
- Running training workflows with custom models
152+
153+
To work through the example, follow the same process as before:
154+
155+
```bash
156+
cd neural-graphics-model-gym-examples
157+
source nb-env/bin/activate
158+
jupyter lab
159+
```
160+
161+
Navigate to `tutorials/nss/custom_model_example.ipynb` and step through the cells.
162+
163+
## Wrapping up
164+
165+
You've now learned how to extend Model Gym beyond NSS with your own models and datasets. This opens up possibilities for experimenting with different neural graphics techniques: denoising, frame interpolation, or custom upscaling approaches tailored to your content.
166+
167+
For more information on model registration, dataset integration, and use case development, see the [Model Gym GitHub repository](https://github.com/arm/neural-graphics-model-gym).
168+
169+
Through this Learning Path, you’ve learned what neural graphics is and why it matters for game performance. You’ve stepped through the process of training and evaluating a model using PyTorch and the Model Gym, and seen how to export that model into VGF (.vgf) for real-time deployment. You’ve also explored how to visualize and inspect the model’s structure using Model Explorer. You can now explore the Model Training Gym repository for deeper integration and to keep building your skills.

0 commit comments

Comments
 (0)