litert-community/HSEmotion-B0-LiteRT

HSEmotion (EfficientNet-B0) — facial emotion recognition on LiteRT GPU

Recognize the 8 AffectNet emotions — anger, contempt, disgust, fear, happiness, neutral, sadness, surprise — from a cropped face, with the whole network running on the LiteRT CompiledModel GPU delegate (no CPU fallback). HSEmotion (EmotiEffLib, Apache-2.0) is an EfficientNet-B0 fine-tuned on AffectNet.

• Architecture: EfficientNet-B0 (tf_efficientnet_b0), 224×224, 8-way head.
• Weights: av-savchenko/face-emotion-recognition enet_b0_8_best_afew.pt · Apache-2.0.
• Size: 8 MB (fp16).

Detected face + emotion distribution on a Pixel 8a; the classifier runs on the GPU.

I/O

• Input: [1, 3, 224, 224] NCHW, RGB, ImageNet-normalized (a cropped face).
• Output: [1, 8] logits, index order: anger, contempt, disgust, fear, happiness, neutral, sadness, surprise.

GPU conversion

Two hurdles, both fixed in build_hsemotion.py:

1. Old-timm pickle. The released weights are a pickled model built with an old timm whose forward is broken under current timm (missing conv_s2d). The state dict is lifted into a fresh timm tf_efficientnet_b0 (num_classes=8, remapping classifier.0.* → classifier.*), which has a working forward — 358/360 tensors match by name and shape, the rest is the remapped classifier.
2. fp16 SqueezeExcite mean → NaN. ⭐The SE block's global mean x.mean((2,3)) over the 112×112 stem map is a single fp16 reduction whose partial sum overflows 65504 → the GPU delegate emits an all-NaN output (it computes the reduction in fp16 even for an fp32 graph; desktop fp16 CPU is exact). Replaced by a hierarchical mean — repeated avg_pool2d over equal-size tiling windows (≤ 49 elements each) — mathematically identical but fp16-safe.

Pixel 8a: 342/342 nodes on the GPU delegate, 1 partition, ~2 ms/inference (fp16). Device fp16 top-1 matches desktop fp32 (logits corr 0.99997); desktop fp16 CPU corr vs PyTorch is 1.0.

Minimal usage

Kotlin (Android, LiteRT CompiledModel GPU)

val model = CompiledModel.create(context.assets, "hsemotion_b0_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()

inputs[0].writeFloat(faceNchw)     // [1,3,224,224] cropped face, ImageNet-normalized
model.run(inputs, outputs)
val logits = outputs[0].readFloat()   // [8] -> softmax + argmax

Python (LiteRT CompiledModel API)

import numpy as np
from ai_edge_litert.compiled_model import CompiledModel

model = CompiledModel.from_file("hsemotion_b0_fp16.tflite")
inputs = model.create_input_buffers(0)
outputs = model.create_output_buffers(0)
inputs[0].write(np.ascontiguousarray(face, np.float32))   # [1,3,224,224]
model.run_by_index(0, inputs, outputs)
logits = outputs[0].read(8, np.float32)                   # argmax -> emotion

The model expects a tightly cropped face (use any face detector; the Android sample uses the built-in android.media.FaceDetector).

License

Apache-2.0 (HSEmotion / EmotiEffLib). Converted with litert-torch.