|
| 1 | +# Example 4 - Batching and Higher-Dimensional Inference with PyTorch |
| 2 | + |
| 3 | +This example demonstrates how to use a PyTorch model trained on 1D vectors to perform |
| 4 | +inference on batched and higher-dimensional data. We start with a simple batch and then |
| 5 | +show how to handle more complex cases, such as time series or multi-dimensional data. |
| 6 | + |
| 7 | +## Description |
| 8 | +A Python file `batchingnet.py` is provided that defines a simple PyTorch 'BatchingNet' |
| 9 | +that takes an input vector of length 5 and applies a single `Linear` layer to multiply |
| 10 | +each input feature by a different value (0, 1, 2, 3, 4). The demo script |
| 11 | +`batching_demo.py` shows how to use this model for unbatched, batched, and |
| 12 | +higher-dimensional inference, illustrating the effect of batching and the location of |
| 13 | +the batch dimension. All outputs are as described in the "Why Batching?" section below. |
| 14 | + |
| 15 | + |
| 16 | +## Dependencies |
| 17 | +To run this example requires: |
| 18 | + |
| 19 | +- CMake |
| 20 | +- Fortran compiler |
| 21 | +- FTorch (installed as described in main package) |
| 22 | +- Python 3 |
| 23 | + |
| 24 | + |
| 25 | +## Why Batching? |
| 26 | +Batching allows you to efficiently process multiple inputs at once, leveraging vectorised |
| 27 | +operations and hardware acceleration. PyTorch modules are designed to handle batched |
| 28 | +inputs by default, as long as the input's last dimension matches the expected feature size. |
| 29 | +The outputs for each case are shown below so you can check your results. |
| 30 | + |
| 31 | +### Unbatched Inference |
| 32 | +Suppose you have a model trained on input vectors of shape `[5]`. You can run inference |
| 33 | +on a single vector as follows: |
| 34 | + |
| 35 | +``` |
| 36 | +input_tensor = torch.ones(5) # shape: [5] |
| 37 | +output = model(input_tensor) |
| 38 | +``` |
| 39 | +Input: |
| 40 | +``` |
| 41 | +tensor([1., 1., 1., 1., 1.]) |
| 42 | +``` |
| 43 | +Expected output: |
| 44 | +``` |
| 45 | +tensor([0., 1., 2., 3., 4.]) |
| 46 | +``` |
| 47 | + |
| 48 | + |
| 49 | +### Simple Batch Inference |
| 50 | +You can run inference on a batch of such vectors by stacking them into a tensor of |
| 51 | +shape `[batch_size, 5]`. |
| 52 | + |
| 53 | +``` |
| 54 | +input_tensor = torch.stack([ |
| 55 | + torch.ones(5), |
| 56 | + 2 * torch.ones(5), |
| 57 | +]) |
| 58 | +output = model(input_tensor) |
| 59 | +``` |
| 60 | +Input: |
| 61 | +``` |
| 62 | +tensor([[1., 1., 1., 1., 1.], |
| 63 | + [2., 2., 2., 2., 2.]]) |
| 64 | +``` |
| 65 | +Expected output: |
| 66 | +``` |
| 67 | +tensor([[0., 1., 2., 3., 4.], |
| 68 | + [0., 2., 4., 6., 8.]]) |
| 69 | +``` |
| 70 | + |
| 71 | +The model will process each row independently. |
| 72 | + |
| 73 | + |
| 74 | +### Higher-Dimensional Inference (e.g., [batch, time, 5]) |
| 75 | +You can also run inference on data with more dimensions, such as `[batch, time, 5]`. |
| 76 | +PyTorch modules like `nn.Linear` operate on the last dimension, so this works as long |
| 77 | +as the last dimension matches the model's input size. |
| 78 | + |
| 79 | +For example, using a tensor with shape `[2, 3, 5]`: |
| 80 | + |
| 81 | +``` |
| 82 | +input_tensor = torch.stack([ |
| 83 | + torch.stack([ |
| 84 | + torch.ones(5), |
| 85 | + 2 * torch.ones(5), |
| 86 | + 3 * torch.ones(5), |
| 87 | + ]), |
| 88 | + torch.stack([ |
| 89 | + 10 * torch.ones(5), |
| 90 | + 20 * torch.ones(5), |
| 91 | + 30 * torch.ones(5), |
| 92 | + ]), |
| 93 | +]) |
| 94 | +output = model(input_tensor) |
| 95 | +``` |
| 96 | +Input: |
| 97 | +``` |
| 98 | +tensor([[[ 1., 1., 1., 1., 1.], |
| 99 | + [ 2., 2., 2., 2., 2.], |
| 100 | + [ 3., 3., 3., 3., 3.]], |
| 101 | +
|
| 102 | + [[10., 10., 10., 10., 10.], |
| 103 | + [20., 20., 20., 20., 20.], |
| 104 | + [30., 30., 30., 30., 30.]]]) |
| 105 | +``` |
| 106 | +Expected output: |
| 107 | +``` |
| 108 | +tensor([[[ 0., 1., 2., 3., 4.], |
| 109 | + [ 0., 2., 4., 6., 8.], |
| 110 | + [ 0., 3., 6., 9., 12.]], |
| 111 | +
|
| 112 | + [[ 0., 10., 20., 30., 40.], |
| 113 | + [ 0., 20., 40., 60., 80.], |
| 114 | + [ 0., 30., 60., 90., 120.]]]) |
| 115 | +``` |
| 116 | + |
| 117 | + |
| 118 | + |
| 119 | +## Key Considerations |
| 120 | +- **Shape:** The last dimension of the input must match the model's expected input |
| 121 | + size (here, 5) otherwise it will result in an error. |
| 122 | +- **Location of Batching Dimension:** PyTorch modules such as `nn.Linear` always operate |
| 123 | + on the last dimension of the input tensor. Any number of leading (batching) dimensions |
| 124 | + are supported, and the operation is applied independently across all of them. |
| 125 | + For example, `[batch, 5]`, `[batch, time, 5]`, or even `[batch1, batch2, ..., 5]` are |
| 126 | + all valid as long as the final dimension is consistent with the feature dimension. |
| 127 | +- **Contiguity:** If you slice or permute tensors, ensure they are contiguous in memory |
| 128 | + before passing to the model (use `.contiguous()` if needed). |
| 129 | + |
| 130 | + |
| 131 | +## Running |
| 132 | + |
| 133 | +To run this example, first install FTorch as described in the main documentation. Then, from this directory, create a virtual environment and install the necessary Python modules: |
| 134 | +``` |
| 135 | +python3 -m venv venv |
| 136 | +source venv/bin/activate |
| 137 | +pip install -r requirements.txt |
| 138 | +``` |
| 139 | + |
| 140 | +You can see how you would perform batching in PyTorch by running the Python batching |
| 141 | +demo: |
| 142 | +``` |
| 143 | +python3 batching_demo.py |
| 144 | +``` |
| 145 | +This will run the BatchingNet model with various inputs and print the expected outputs |
| 146 | +as described above. |
| 147 | + |
| 148 | +To save the BatchingNet model to TorchScript for use in Fortran, run: |
| 149 | +``` |
| 150 | +python3 pt2ts.py |
| 151 | +``` |
| 152 | +This will generate `saved_batchingnet_model_cpu.pt` in the current directory. |
| 153 | + |
| 154 | +At this point you no longer require Python, so can deactivate the virtual environment: |
| 155 | +``` |
| 156 | +deactivate |
| 157 | +``` |
| 158 | + |
| 159 | +To call the saved BatchingNet model from Fortran, repeating the different batching |
| 160 | +approaches in the Python demo compile the `batchingnet_infer_fortran.f90` file. |
| 161 | +This can be done using the included `CMakeLists.txt` as follows: |
| 162 | +``` |
| 163 | +mkdir build |
| 164 | +cd build |
| 165 | +cmake .. -DCMAKE_PREFIX_PATH=<path/to/your/installation/of/library/> -DCMAKE_BUILD_TYPE=Release |
| 166 | +cmake --build . |
| 167 | +``` |
| 168 | +(Note that the Fortran compiler can be chosen explicitly with the `-DCMAKE_Fortran_COMPILER` flag, and should match the compiler that was used to locally build FTorch.) |
| 169 | + |
| 170 | +To run the compiled code calling the saved BatchingNet TorchScript from Fortran, run the executable with an argument of the saved model file: |
| 171 | +``` |
| 172 | +./batchingnet_infer_fortran ../saved_batchingnet_model_cpu.pt |
| 173 | +``` |
| 174 | + |
| 175 | +This runs the model with single, batched, and multidimensional batched input arrays and should produce the output: |
| 176 | +``` |
| 177 | +--- Single input output: --- |
| 178 | +0.00000000 1.00000000 2.00000000 3.00000000 4.00000000 |
| 179 | +
|
| 180 | +---Batched input output: --- |
| 181 | + 0.00000000 1.00000000 2.00000000 3.00000000 4.00000000 |
| 182 | + 0.00000000 2.00000000 4.00000000 6.00000000 8.00000000 |
| 183 | +
|
| 184 | +--- Multidimensional batched input/output: --- |
| 185 | +Input (1,1): 1.00000000 1.00000000 1.00000000 1.00000000 1.00000000 |
| 186 | +Output (1,1): 0.00000000 1.00000000 2.00000000 3.00000000 4.00000000 |
| 187 | +Input (2,3): 30.0000000 30.0000000 30.0000000 30.0000000 30.0000000 |
| 188 | +Output (2,3): 0.00000000 30.0000000 60.0000000 90.0000000 120.000000 |
| 189 | +
|
| 190 | +BatchingNet Fortran example ran successfully |
| 191 | +``` |
| 192 | + |
| 193 | +--- |
| 194 | + |
| 195 | +## Further Exploration |
| 196 | + |
| 197 | +- incorrect Dimension order |
| 198 | + - A commented-out line in the Fortran code demonstrates how an incorrect feature |
| 199 | + dimension (e.g., `[2, 4]`) will cause a runtime error. |
| 200 | + ```fortran |
| 201 | + call torch_model_forward(model, in_tensors_bad, out_tensors_bad) |
| 202 | + ``` |
| 203 | + - Uncomment, rebuild and re-run to observe the error message for incorrect |
| 204 | + batching layout. |
| 205 | +- Try using different batch and time dimensions. |
| 206 | +- Change the nature of the model and observe the effect of batching. |
0 commit comments