|
| 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