File size: 3,563 Bytes
9705b2b 98a233a 9705b2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | ---
license: bsd-3-clause
library_name: LiteRT
pipeline_tag: image-to-image
tags: [litert, tflite, on-device, android, gpu, style-transfer, neural-style, fast-neural-style]
base_model: pytorch/examples
---
# Fast Neural Style Transfer — LiteRT (on-device, fully-GPU, 4 styles)
Fast neural **style transfer** ([PyTorch examples](https://github.com/pytorch/examples/tree/main/fast_neural_style)
`TransformerNet`, Johnson et al.), converted to **LiteRT** and running **fully on the `CompiledModel` GPU**
(ML Drift) on Android. Applies an artistic style to a photo — **4 styles** (candy / mosaic / rain_princess /
udnie), each a **3.5 MB** fp16 graph.
## On-device (Pixel 8a, Tensor G3 — verified)
| | |
|---|---|
| nodes on GPU | **350 / 350** LITERT_CL (full residency) |
| inference | **~9 ms** (256×256) |
| size | 3.5 MB per style (fp16) |
| accuracy | device-vs-PyTorch corr **0.9998–0.9999** (all 4 styles) |
```
image[1,3,256,256] (RGB 0-255) →[GPU: TransformerNet]→ stylized[1,3,256,256] (RGB 0-255)
```
## Minimal usage
**Android (Kotlin, CompiledModel GPU)**
```kotlin
val model = CompiledModel.create(context.assets, "style_candy_fp16.tflite",
CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
inputs[0].writeFloat(chw) // [1,3,256,256] RGB 0-255, NCHW
model.run(inputs, outputs)
val stylized = outputs[0].readFloat() // [1,3,256,256] RGB 0-255
```
**Python (desktop verification)**
```python
import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter
img = Image.open("photo.jpg").convert("RGB")
w, h = img.size; s = min(w, h)
img = img.crop(((w-s)//2, (h-s)//2, (w+s)//2, (h+s)//2)).resize((256, 256))
x = np.asarray(img, np.float32).transpose(2, 0, 1)[None] # 0-255, no normalization
# candy / mosaic / rain_princess / udnie
it = Interpreter(model_path="style_candy_fp16.tflite"); it.allocate_tensors()
it.set_tensor(it.get_input_details()[0]["index"], x); it.invoke()
y = it.get_tensor(it.get_output_details()[0]["index"])[0] # [3,256,256] RGB 0-255
Image.fromarray(y.transpose(1, 2, 0).clip(0, 255).astype(np.uint8)).save("stylized.png")
```
## How it converts (litert-torch) — three numerically-exact re-authorings
1. **`ReflectionPad2d` → zero-pad** (`GATHER_ND` → `PAD`; border-only difference).
2. **Large conv activations → conv-weight scaling.** The conv outputs reach ≈ |5000|, where the Mali delegate's
fp16 conv accumulation loses precision → garbage (device corr 0.34 at full residency — *residency ≠
correctness*). Each conv is followed by an `InstanceNorm` (which is **scale-invariant**), so scaling those
conv weights down so the output is ≈ |10| is **exact** (IN output unchanged) and keeps the fp16 accumulation
precise → corr 1.0.
3. **`InstanceNorm` → SafeInstanceNorm** (down-scaled-domain spatial reduction, fp16-safe; SafeLayerNorm class).
Upsample is `interpolate(nearest)` (no transposed conv → no ZeroStuff). Result: banned ops NONE, ≤4D,
tflite-vs-torch corr **1.0**, device-vs-torch corr **0.9999**.
## Preprocessing
Center-crop to square, resize to 256×256, RGB **0–255** (no normalization), NCHW. Output is 0–255 RGB (clamp).
## License
[BSD-3-Clause](https://github.com/pytorch/examples/blob/main/LICENSE). Upstream:
[pytorch/examples](https://github.com/pytorch/examples/tree/main/fast_neural_style).
|