Tristan Penman's Blog

Robotics in Go: Building a Raspberry Pi Rover

Fri, 22 May 2026 00:00:00 +0000

When we consider programming languages for robotics, those that come to mind often include established players such as C and C++, Python, or even those specifically designed for scientific and numerical computing, such as MATLAB. In this post, we look at Go as a viable programming language for robotics.

This post builds upon content I presented at the Melbourne Go Meetup on May 14, 2026. My talk centred around a project called Rover Kit, a small rover that uses Go for motor control and sonar readings. The Rover Kit source code has been published on GitHub. The Rover Kit v0.1.0 release contains the version of the code discussed below.

Introduction
Bare Metal Robotics
Hardware Interfaces
- GPIO
- UART
- SPI
- I²C
Go Libraries
- Gobot
- Periph
- Serial
Rover Kit
Motor Control
Sonar Reader
Web Bridge
- WebSocket Server
- Gamepad Support
Next Steps
References

Introduction

So much of the appeal of robotics is its breadth as a subject. It touches upon mathematics, mechanical and electrical engineering, and (of course) software. Just as broad are the many different kinds of robots we can build. The most ambitious include bipedal humanoid robots, actuated robots that operate in factories, or even surgical devices.

The Hardware

This post focuses on a small rover that includes four DC motors (driven by an Adafruit Motor HAT) and several HC-SR04 ultrasonic distance sensors. The chassis and wheels are from a Whippersnapper Runt Rover kit. With everything connected, the rover looks something like this:

Hidden within the chassis is a 12V battery pack to power the DC motors.

It also includes a 5V USB power supply that powers a Raspberry Pi Zero W - this is the brains of the operation:

Although fairly modest as far as robotics is concerned, this kit allows us to explore the software stack while steering clear of advanced topics (e.g. inverse kinematics and frames of reference).

Robotics Fundamentals

A typical robot incorporates actuation (wheels, motors, or other moving parts), sensors to perceive the environment, and some degree of planning and control. Tying this all together are software systems that ensure the robot achieves the task it was designed to do.

The demands placed on robotics software are non-trivial. They include:

Real-time control - operating within precise timing constraints
Embedded environments - low power and harsh operating conditions (heat, dust, moisture)
Probabilistic planning - working with sensor data that might be noisy or incomplete
Safety systems and critical infrastructure - ensuring safe and correct operation, under a wide range of conditions

To address these concerns, large-scale hardware projects will often turn to a mature robotics framework that offers component isolation, message passing, device drivers, diagnostics, and deployment tooling.

A popular choice today is ROS, otherwise known as the Robot Operating System.

Robot Operating System

ROS provides middleware, standard message types, navigation and perception packages, and tools for visualisation and debugging. It also offers broad device support, which makes it appealing for research and prototyping.

ROS has historically had two major release lines. ROS 1 was built around a central roscore process, but reached end-of-life in 2025. ROS 2 moved away from that central control node architecture and added features aimed at commercial and industrial robotics, including QoS, device discovery, and improved security.

Message Passing

A common theme, though, is message passing. ROS 1 and ROS 2 are both designed around message passing as a core concept. Individual components can be responsible for just one aspect of a robot’s behaviour, with components coordinating towards larger goals by publishing and subscribing to different topics on a message bus:

Although the message passing concept has proven scalable and robust, ROS itself has a steep learning curve. It also comes with resource demands that may strain small devices. This complexity may distract from the actual goal of a hobbyist: having fun with a project that combines hardware and software. For that we want something more akin to “bare metal robotics”.

Bare Metal Robotics

Bare metal robotics captures the notion of doing robotics with as little overhead as possible.

Python is a common choice for hobbyist robotics, thanks in part to its simplicity. Python features a wide range of libraries for GPIO communication and other hardware interfaces. However, a major drawback is performance, as it suffers from being an interpreted language. Deployment can also be challenging, due to the number of dependencies that must be installed on the target device.

Go for Robotics

It turns out that Go is also a great option for bare metal robotics.

While it may not be perfect, Go has several properties that make it an appealing choice:

Concurrency model - goroutines and channels offer a clean mapping to asynchronous sensor/actuator loops
Memory safety - no pointer arithmetic and a memory-safe runtime make it safer than C++ (for many classes of bugs… in theory)
Simplicity & readability - Go is often easier to read than C++ or Python
Cross-compilation - often produces self-contained binaries with minimal runtime dependencies
Iteration - fast builds, simple dependency management, good testing support

The concurrency model and support for cross-compilation are two of the most interesting properties (in my opinion!).

Concurrency Model

Robots combine multiple activities that are conceptually independent, but must coordinate with one another to achieve various tasks. A simple robot may combine sensors, actuators (e.g. motor controllers) and basic navigation. Each of these components (or processes) needs to run concurrently.

Go’s concurrency model maps nicely to that problem. A goroutine can own a long-running activity, such as sampling a sonar sensor once per second, while a channel becomes the handoff point for the values it produces. In turn, Go’s treatment of channels as first-class language features leads to several nice concurrency patterns.

Go’s context package also provides clean shutdown semantics. Commands can create a context tied to SIGINT or SIGTERM, pass it down into individual goroutines, and block until cancellation.

Cross-Compilation

Cross-compilation is another major benefit! The Go compiler makes it extremely easy to cross-compile binaries that run natively on the Raspberry Pi Zero W.

The compiler provides built-in support for ARMv6, which is the instruction set used by the Pi Zero. Cross-compilation is as simple as running go build on your host machine with the appropriate command line arguments and environment variables:

env GOOS=linux GOARCH=arm GOARM=6 \
  go build -ldflags "-w" -o bin/motor-control cmd/motor-control/main.go

Go Drawbacks

The primary drawback of Go is the garbage collector (GC), which can lead to unpredictable latency spikes. Timing discrepancies can impact perception and coordination. For certain motor configurations, this could lead to overshooting or undershooting actuation targets.

While it’s possible to configure garbage collection, or even switch it off entirely, these are beyond the scope of this post.

Hardware Interfaces

Embedded devices support a range of hardware I/O interfaces, each designed to accommodate different hardware design goals. A common characteristic of the interfaces described below is that they are all serial.

Readers who are old enough may recall owning a computer that had both serial and parallel ports:

Image credit: Difference between Serial Port and Parallel Port by GeeksForGeeks

An obvious difference between the two ports is the number of pins required. The parallel port has a wider bus, which means it can transmit more bits at the same time, while serial is designed to transmit one bit at a time.

Serial may seem limiting, but it forms the foundation of various hardware interfaces. This includes the ubiquitous USB (Universal Serial Bus) and the GPIO pins commonly found on embedded devices.

GPIO

GPIO (or General-Purpose Input/Output) refers to a set of digital pins that can be controlled via software. In hobby electronics projects, it is common to wire up GPIO pins to LEDs, sensors, and other small components. Each SBC has its own GPIO pin configuration, such as this one for the Raspberry Pi:

Image credit: How to Use Raspberry Pi GPIO Pins – Python Tutorial by Seeed Studio.

Groups of GPIO pins can be used to implement more sophisticated protocols or hardware interfaces. We’ll introduce some of those now.

UART

One of the most common interfaces you’ll find on embedded devices is UART, short for Universal Asynchronous Receiver/Transmitter. UART is a protocol for serial communication between two devices, using separate transmit and receive signal wires plus a shared ground. Two UART devices are connected via cross-over wiring:

For signalling purposes, each character is preceded by a start bit (always a zero) and ends with one or more stop bits (logical ones). The signal line is held in the ‘one’ state when nothing is being transmitted:

Standard UART data rates range from 9600 bits per second (bps) up to 115200 bps. High-speed UART implementations exist, but are less common. A more common alternative is USART, short for Universal Synchronous/Asynchronous Receiver-Transmitter, which includes an additional clock wire for synchronisation, enabling faster transfer rates.

SPI

SPI (Serial Peripheral Interface) is a fast synchronous serial bus commonly used to talk to chips on the same board. In a typical setup, there is one SPI controller, which is responsible for communicating with one or more SPI peripherals. Peripheral chips are each targeted using a dedicated chip select pin, often labelled CS or SS.

The following diagram shows three SPI peripherals wired up to a single SPI controller. Note the presence of three CS/SS pins:

Examples of SPI peripherals include displays, sensors, and storage devices (think SD cards). When working with embedded devices you may encounter SPI Flash, a kind of non-volatile flash memory that is capable of retaining data without power. SPI Flash is well suited to storing data and program code that needs to be loaded when a device is first powered up.

I²C

I²C (Inter-Integrated Circuit) is a two-wire synchronous serial bus designed for connecting multiple chips via the same circuit.

Unlike SPI, I²C is an addressed bus, which means that individual chips are targeted via an address encoded into the bitstream. Chips are connected to the bus via two wires: serial data (SDA) and serial clock (SCL). It is assumed that there is an I²C controller which controls the bus and clocks data from target devices when reading.

A start of a frame consists of a 7-bit address, a read/write bit which indicates whether the controller is sending or receiving data, and an ACK bit. The control then sends data, or receives data, depending on the R/W bit.

The Adafruit Motor HAT in this project uses an I²C-based PWM driver chip. This means that, with just two signal wires connected to the Raspberry Pi, we can control four DC motors!

Go Libraries

We’ve seen above that Go can be used to target platforms such as the Raspberry Pi, BeagleBone or Jetson Nano, and we’ve introduced some of the hardware interfaces available on these devices. The next question is how we can make use of those hardware interfaces.

This section will look at three in detail: Gobot, Periph and Serial. These are the three libraries used in Rover Kit.

There are many others we could look at (such as gpiocdev which gets an honourable mention), but unfortunately we don’t have time to cover them all.

Gobot

Gobot is a robotics framework for Go that provides APIs for devices, sensors, and actuators across a wide range of platforms. Using the provided drivers and adaptors, you can write higher-level behaviour once and retarget from one board or device family to another with minimal changes.

The following example uses the firmata adaptor to control an LED connected to an Arduino:

package main

import (
  ...
)

func main() {
  firmataAdaptor := firmata.NewAdaptor("/dev/ttyACM0")
  led := gpio.NewLedDriver(firmataAdaptor, "13")

  work := func() {
    gobot.Every(1*time.Second, func() {
      led.Toggle()
    })
  }

  robot := gobot.NewRobot("bot",
    []gobot.Connection{firmataAdaptor},
    []gobot.Device{led},
    work,
  )

  // Start the main event loop
  robot.Start()
}

Gobot provides adaptors for a wide range of SBC devices, including the Raspberry Pi, BeagleBone, and Tinker Board. It also supports drones and robotics kits.

Periph

Periph is a low-level hardware I/O library for Go. It provides direct access to GPIO, I²C, SPI, PWM, and other interfaces commonly found on SBCs. It is especially useful when you want tight control of peripheral communication without needing a heavier robotics framework.

We will see later that Periph can be used to drive motors via the Adafruit Motor HAT, and measure distance using the HC-SR04 ultrasonic sensors.

Serial

The last library I want to mention is go.bug.st/serial, a cross-platform serial port library for Go.

As you can see, this library makes it easy to open and read from a serial device:

// configure
mode := &serial.Mode{
  BaudRate: 115200,
  DataBits: 8,
  StopBits: serial.OneStopBit,
  Parity:   serial.NoParity,
}

// open port
port, err := serial.Open("/dev/ttyACM0", mode)
if err != nil {
  return nil, err
}

// read bytes from UART connection
buff := make([]byte, 128)
bytesRead, err := port.Read(buff)
if err != nil {
  return nil, err
}
_ = bytesRead

This can be used to delegate reading from ultrasonic sensors to a separate microcontroller, which is something we’ll explore in a future post!

Rover Kit

In this section, we take a closer look at the hobby project that inspired this post: Rover Kit. The project adopts a modular architecture, with separate Go executables to drive motors, read sonar sensors, and send commands via a web interface.

MQTT provides pub/sub messaging between components:

MQTT is a standard protocol for communication between IoT devices. This is similar to the message bus we saw in ROS, and has the advantage of being highly portable. At the time of writing, there are over 20 server implementations and even more client libraries.

MQTT Client

We make use of the Eclipse Paho MQTT Go client to connect to an MQTT broker and pass messages between components.

Basic connection:

client := mqtt.NewClient(opts)
connectToken := client.Connect()
connectToken.Wait()
if err := connectToken.Error(); err != nil {
  log.Fatalf("failed to connect to broker=%s err=%v", brokerURL, err)
}

Subscription example:

opts.SetOnConnectHandler(func(client mqtt.Client) {
  log.Printf("connected to broker=%s", brokerURL)
  token := client.Subscribe(motorCmdTopic, 1, subscriber(ctx, driver, gate))
  token.Wait()
  if err := token.Error(); err != nil {
    log.Printf("failed to subscribe topic=%s err=%v", motorCmdTopic, err)
    return
  }
  log.Printf("subscribed topic=%s", motorCmdTopic)
})

Publishing example:

s.mqttClient.Publish(defaultMotorCmdTopic, 0, false, message)

Peripherals

Recall that our rover hardware included:

Four DC motors
Four HC-SR04 ultrasonic distance sensors

The motor HAT is controlled over I²C through the Raspberry Pi’s GPIO header, while the ultrasonic sensors use ordinary GPIO pins. Rover Kit provides motor driver implementations for Gobot and Periph. Supporting both is not strictly necessary for such a simple project. However, it serves some pedagogical value by allowing us to compare the implementations side-by-side.

Project Structure

Rover Kit follows a standard Go project structure. The three commands live in cmd/, while supporting code lives in pkg/. There is no internal/ directory, as the project is intended to be completely open.

You will also find configuration files for deployment on the Raspberry Pi:

deploy/ - systemd configuration files
mqtt/ - mosquitto configuration files

A Makefile has also been included, making it very easy to cross-compile for ARMv6.

make

Next, we’ll look at the three commands: motor-control, sonar-reader, and web-bridge.

Motor Control

The motor-control command knows how to interact with our DC motors, via the Adafruit Motor HAT. This is backed by the Driver interface, which declares the main behaviours that we need to support for our four DC motors:

type Driver interface {
	Forwards(ctx context.Context) error
	Backwards(ctx context.Context) error
	SpinCW(ctx context.Context) error
	SpinCCW(ctx context.Context) error
	Stop(ctx context.Context) error
	Close(ctx context.Context) error
}

Pulse Width Modulation

Pulse Width Modulation (or PWM) controls a DC motor by rapidly switching the motor’s power supply on and off instead of varying the voltage directly.

You can picture this as a cycle. During the “on” part of each cycle, the motor receives 100% voltage. During the “off” part of the cycle, it receives no voltage. Because this happens so quickly, the motor behaves as if it is receiving an average voltage. The proportion of “on” vs “off” time is called the “duty cycle”. A 20% duty cycle means the motor is powered 20% of the time, and unpowered for the remaining 80%.

The Adafruit Motor HAT handles this for us. However, we still need to tell it what the desired duty cycle is for each motor. This is communicated via I²C.

Gobot Driver

Gobot provides an adaptor for the Raspberry Pi and an I²C driver for the Adafruit Motor HAT:

// Create an adaptor for the Raspberry Pi I2C interface
adaptor := raspi.NewAdaptor()
if err := adaptor.Connect(); err != nil {
  return nil, fmt.Errorf("connect raspi adaptor: %w", err)
}

// Initialise the motor hat driver
hat := i2c.NewAdafruitMotorHatDriver(
  adaptor,
  i2c.WithBus(1),
  i2c.WithAddress(defaultGobotMotorHatI2C),
)

// Start the driver
if err := hat.Start(); err != nil {
  _ = adaptor.Finalize()
  return nil, fmt.Errorf("start Adafruit motor hat driver: %w", err)
}

Although Gobot also offers a high-level event loop, we won’t use it for such a simple project. We simply send commands to the motor HAT, using a setMotor helper function:

func (d *GobotDriver) setMotor(motor int, throttle float64) error {
	t := math.Max(-1, math.Min(1, throttle))
	speed := int32(math.Round(math.Abs(t) * 255))

	if err := d.hat.SetDCMotorSpeed(motor, speed); err != nil {
		return err
	}

	switch {
	case t > 0:
		return d.hat.RunDCMotor(motor, i2c.AdafruitForward)
	case t < 0:
		return d.hat.RunDCMotor(motor, i2c.AdafruitBackward)
	default:
		return d.hat.RunDCMotor(motor, i2c.AdafruitRelease)
	}
}

Using the helper simplifies the implementation of individual actions, as we can see in Forwards below:

func (d *GobotDriver) Forwards(context.Context) error {
	d.mu.Lock()
	defer d.mu.Unlock()

	values := [4]float64{-1, 1, -1, 1}
	for i, value := range values {
		motor := i + 1
		if err := d.setMotor(motor, value); err != nil {
			return fmt.Errorf("forwards motor %d: %w", motor, err)
		}
	}
	return nil
}

Periph Driver

Another option is Periph, which also happens to support the PCA9685 chip on the Adafruit Motor HAT.

Unlike Gobot’s Motor HAT driver, we must manually specify the chip’s PWM motor configuration:

var periphMotorChannels = [4]periphMotorChannel{
	{pwm: 8, in1: 10, in2: 9},
	{pwm: 13, in1: 11, in2: 12},
	{pwm: 2, in1: 4, in2: 3},
	{pwm: 7, in1: 5, in2: 6},
}

The equivalent setMotor helper is a little more involved. Some error handling has been omitted for brevity:

func (d *PeriphDriver) setMotor(motor int, throttle float64) error {
	if motor < 1 || motor > len(periphMotorChannels) {
		return fmt.Errorf("unsupported motor index %d", motor)
	}
	ch := periphMotorChannels[motor-1]

	t := math.Max(-1, math.Min(1, throttle))
	duty := pwmDutyFromThrottle(t)

	switch {
	case t > 0:
		d.setChannel(ch.in1, true)
		d.setChannel(ch.in2, false)
	case t < 0:
		d.setChannel(ch.in1, false)
		d.setChannel(ch.in2, true)
	default:
		d.setChannel(ch.in1, false)
		d.setChannel(ch.in2, false)
	}

	if duty == 0 {
		return d.pca.SetFullOff(ch.pwm)
	}

	return d.pca.SetPwm(ch.pwm, 0, duty)
}

Sonar Reader

The sonar-reader command is responsible for taking readings from each of the HC-SR04 ultrasonic sensors. This is backed by the Provider interface:

type Reading struct {
	DistanceCM float64
	DurationUS float64
	Timestamp  time.Time
}

type Provider interface {
	Open(ctx context.Context) chan Reading
	Close(ctx context.Context) error
}

Sonar Sensors

Interacting with the HC-SR04 sensors is a little tricky, due to their timing sensitivity. Each sensor component has four pins, a transmitter (T / TX), and a receiver (R / RX):

To measure the distance to the nearest obstacle, the transmitter first emits an ultrasonic pulse, then the receiver detects the echo of that pulse. The time between these two events is used to calculate the distance.

We can tell the sensor to emit a pulse by holding the trigger pin high for a short period. The sensor then raises the echo pin for the duration of a reflected ultrasonic pulse. Measuring that echo duration gives the round-trip travel time. We multiply by the speed of sound to calculate the total distance, then divide by two since we only want the distance in one direction:

Voltage Divider

It’s also worth noting voltage requirements. The HC-SR04 accepts 3.3V or 5V on the trigger pin, but the echo pin outputs 5V. However, the GPIO pins on the Raspberry Pi are not 5V tolerant. This means we must use a voltage divider to drop the voltage on the echo pin to 3.3V.

Voltage divider circuits are a dime-a-dozen online, and you’ll want to choose one that suits your particular hardware design. Here’s a photo from when I was halfway through soldering the small voltage divider board included in the rover:

This iteration of the board supported two HC-SR04 sensors.
The final board supports four HC-SR04 sensors.

Periph Provider

Using Periph, we can interact with the HC-SR04 via GPIO pins on the Raspberry Pi:

// First we get references to appropriate GPIO pins:
trig := gpioreg.ByName("GPIO18")
echo := gpioreg.ByName("GPIO24")

// Example: Set trigger pin to low
trig.Out(gpio.Low)

// Example: Read from the echo pin
value := echo.In(gpio.Float, gpio.NoEdge)

To take sonar readings indefinitely, we use a for-ever loop inside a goroutine:

for {
  // set trigger pin high
  err := p.trig.Out(gpio.High)

  // sleep for long enough that the sonar will register the signal
  time.Sleep(triggerPulse)

  // set trigger low
  err = p.trig.Out(gpio.Low)

  // wait for echo high
  start := time.Now()
  for p.echo.Read() != gpio.High {
    if time.Since(start) > echoTimeout {
      log.Printf("timed out waiting for echo high")
      return
    }
  }

  // wait for echo low
  start = time.Now()
  for p.echo.Read() != gpio.Low {
    if time.Since(start) > echoTimeout {
      log.Printf("timed out waiting for echo low")
      return
    }
  }

  // calculate distance based on duration
  end := time.Now()
  duration := end.Sub(start)
  distance := duration.Seconds() * soundSpeedCMPerS / 2

  c <- Reading{
    DistanceCM: distance,
    DurationUS: float64(duration.Microseconds()),
    Timestamp:  time.Now(),
  }

  time.Sleep(sampleInterval)
}

UART Provider

Due to the timing-sensitive nature of ultrasonic readings, we also provide a UART provider. This can be used to gather readings from an external microcontroller (such as an STM32-based MCU).

We will look at how this works in the next post.

Web Bridge

Finally, we have the Web Bridge, which allows us to control the rover from a web browser. The web interface is intentionally small. It consists of a single HTML page (index.html), which loads a stylesheet and vanilla JavaScript:

After loading, the web client opens a WebSocket connection back to /ws on the same host. This connection sends rover commands to the server and displays messages received from the server, including distance measurements.

WebSocket Server

The WebSocket endpoint is where the web UI meets MQTT. The web bridge is responsible for validating messages and publishing them via MQTT.

For example, when the browser sends the message {"type":"spin_ccw"}, the web bridge parses and validates the JSON payload. Malformed or unsupported payloads are rejected, and valid payloads are forwarded to the motor command topic (rover/motor/cmd).

The web bridge also subscribes to the sonar sample topic (rover/sonar/sample). When an ultrasonic reading is received, the web bridge publishes it to every connected WebSocket client.

Gamepad Support

The web UI can also translate browser gamepad input into the same command messages used by the on-screen controls. This keeps the backend unchanged: whether a command comes from a button click, a keyboard shortcut, or a gamepad axis, the WebSocket server still receives JSON, validates it, and forwards it to MQTT.

For testing I used an Xbox 360 Gamepad, connected via USB:

Next Steps

The next challenge in this project is to delegate reading from ultrasonic sensors to a separate microcontroller. I will be using an STM32F407 Discovery Board, a common microcontroller development board based on the Arm Cortex-M4 processor. This will be used to send distance readings back to the Raspberry Pi via a USB UART connection.

Stay tuned for the next post, where we’ll use TinyGo to compile and run Go binaries on the STM32.

References

Libraries and frameworks:

Gobot website - A framework for robots, drones, and the Internet of Things
Periph website - A standalone hardware library with limited external dependencies
Eclipse Paho MQTT Go client - Connect to an MQTT message broker from Go

Videos:

Understanding I2C - Great ~10 minute video by Rohde & Schwarz

Embedded Go:

Embedded Go - Bare-metal programming with Go
Embedded Go STM32 - Support for the STM32 microcontrollers
TinyGo - Bring Go to embedded systems and to the modern web using a compiler based on LLVM

Machine Translation at the Edge: Part 2

Fri, 27 Feb 2026 00:00:00 +0000

In the previous post, we introduced the machine translation problem and worked through the fundamentals of the MarianNMT framework. For this installment, we’ve set ourselves a goal to port Marian to the Rockchip NPU. More precisely, we want to download a pretrained MarianMT model from Hugging Face, convert it to Rockchip’s RKNN format, then perform inference directly on the Rockchip NPU - all using Python!

Throughout this journey, we discuss issues such as quantization, input padding and layer compatibility. For those wanting more detail, the complete source code for this project can be found on GitHub. This includes a native C++ implementation too!

Part 1 Recap
Preflight
Converting to ONNX
- Marian-ONNX-Converter
- Model Files
Converting to RKNN
- Quantization
- Dynamic Ops
Inference Engine
Decoding Strategies
- Greedy Decoding
- Beam Search
Benchmarks
Wrap Up
References

Part 1 Recap

When it comes to Machine Translation, there are several well-established models we can choose from. After exploring a few options, we ultimately decided to focus on Marian Transformer-based models. These models have seen adoption in industry and are available for a wide range of language pairs. Marian is also simple enough that if we wanted to support transformer-based models trained using another toolkit, such as Sockeye or OpenNMT, many of the same techniques would apply.

Let’s visit what “Marian” actually means, starting with MarianNMT.

As a quick historical note, Marian NMT was named in honour of Marian Rejewski, a Polish mathematician and cryptologist.

MarianNMT

MarianNMT is a pure C++ machine learning toolkit, specialised for training Neural Machine Translation (NMT) models. MarianNMT can be compiled to run on the CPU, or with GPU-acceleration using CUDA. Although it can be used to train a range of different machine translation models (e.g. LSTMs and Seq2Seq models) it is most commonly used to train Transformer models.

MarianNMT includes an efficient inference engine, also written in C++. This is highly portable, but is not written in a way that enables adoption of the RKNN API. If we tried to cross-compile Marian for Rockchip devices, processing would still happen on the CPU, which is not what we want.

So we’ll need to find another implementation as a reference point. This is where Hugging Face enters the scene.

Hugging Face

Hugging Face is a platform for sharing machine learning models. It’s an easy way to find and access pretrained models, as well as training and inference pipelines. Models published on Hugging Face typically include a “Model Card” containing useful information about the model, how it was trained, and how to use it. You can think of this like the README file you would find in a GitHub repo.

There are a range of Marian-compatible models on Hugging Face. As with the first post, we will focus on an OPUS MT English-to-French model known as Helsinki-NLP/opus-mt-en-fr. This is paired with inference implementation based on Hugging Face Transformers.

Hugging Face Transformers

The Hugging Face Transformers library supports models for text generation, machine translation, summarization and classification (just to name a few). It also provides support for training workflows, model configuration management, and interoperability with backends such as PyTorch and TensorFlow.

To streamline the experience of using a model, Hugging Face provides instructions for various Cloud AI providers and on your own system. Getting started with Helsinki-NLP/opus-mt-en-fr only requires a few lines of code:

Hugging Face Transformers has a built-in Marian implementation, which makes getting started very easy!

MarianNMT vs MarianMT

The last thing to note is the difference between MarianNMT and MarianMT, which are easily confused…

When we say MarianNMT, we’re referring to the C++ toolkit introduced above. When we say MarianMT we’re referring to a Marian implementation based on Hugging Face Transformers. It comes with a set of pretrained models, and this is what we’ll use as the basis for our Rockchip NPU implementation.

Preflight

Let’s get MarianMT up and running, so that we have a baseline for the rest of this project. Note: Everything you see below should be executed on a PC.

The first thing we’ll need to do is install the transformers library, and some of its dependencies:

pip install \
  sentencepiece \
  torch \
  transformers

This command also installs compatible versions of PyTorch and SentencePiece, which are dependencies for MarianMT.

SentencePiece

You may have noticed on the Helsinki-NLP/opus-mt-en-fr model card earlier, that pre-processing includes “normalization + sentencepiece”. This means we’ll need to use Google’s SentencePiece library to tokenize input.

SentencePiece is an unsupervised text tokenizer and detokenizer, based on the byte-pair-encoding (BPE) and unigram algorithms. It is responsible for splitting a source text into subwords. We also use it to convert model output into human-readable sentences.

Test Script

We can put this all together using a simple Python script:

from transformers import MarianMTModel, MarianTokenizer
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

model_name = 'Helsinki-NLP/opus-mt-en-fr'
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name).to(device)

text = ["Hello, I am a fish. What kind of animal are you?"]

inputs = tokenizer(text, return_tensors="pt", padding=True).to(device)
translated = model.generate(**inputs)
translation = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]

print(translation)

We’ve deviated from the model card a bit here, opting to use Marian-specific classes instead of relying on the AutoTokenizer and AutoModelForSeq2SeqLM. This is just for clarity.

The MarianTokenizer and MarianMTModel classes know how to tokenise input and transform, respectively. The tokeniser is also used to convert tokens generated by the model into human-readable output. MarianTokenizer does this using SentencePiece behind the scenes.

When we run the test script above, we see that it successfully translates the English phrase “Hello, I am a fish. What kind of animal are you?” into French:

$ python test.py
["Bonjour, je suis un poisson, quel genre d'animal êtes-vous ?"]

The marian-rknn repo contains a more robust version of this script, called preflight.py.

Note: The first time you run this script it will download the necessary pretrained model files. The model will be stored in a local cache, so that subsequent invocations don’t require another download.

PyTorch vs ONNX

The Hugging Face model downloaded above will be in PyTorch format, since that’s what Hugging Face Transformers uses internally. The question we need to answer is how to go from PyTorch format to RKNN format.

The key challenge is that RKNN Toolkit cannot just run or convert a PyTorch model by executing arbitrary Python code. We have two paths we can follow. The first uses TorchScript, which is a way to serialize and save models to disk so they can be used in alternative runtime environments (e.g. C++). Unfortunately, TorchScript has been deprecated since PyTorch 2.9.

Our second choice is ONNX - the Open Neural Network Exchange format. Once a model has been converted to ONNX, it is relatively easy to convert it to RKNN format. The ONNX runtime also supports the concept of Runtime Execution Providers, making it possible to run a model on a wide range of hardware platforms.

Converting to ONNX

With a little bit of research, we find that PyTorch directly handles ONNX exports. This seems relatively simple at first:

from transformers import MarianMTModel
import torch

model = MarianMTModel.from_pretrained('Helsinki-NLP/opus-mt-en-de')
dummy_input = torch.randint(0, 100, (1, 64))  # token IDs
torch.onnx.export(model, (dummy_input,), "marian.onnx", opset_version=13)

This would work well if you wanted to run the model on a GPU. In practice there’s more work required to export the model in a way that will cleanly convert to RKNN format and run on the Rockchip NPU. Thankfully, Kevin Costa’s Marian-ONNX-Converter project will save us a great deal of time and effort!

Marian-ONNX-Converter

Marian-ONNX-Converter does a few things that will be helpful for us. It exports a MarianMT model into a split encoder/decoder ONNX layout, and preserves the extra artifacts that are required to invoke the model outside of MarianMT. It also provides an ONNX implementation which can run on the CPU, or on a GPU using CUDA.

Marian-ONNX-Converter also happens to support quantization, however we disable quantization when converting from PyTorch to ONNX, as this would introduce layers into the model that are not supported by the RKNN API.

To convert the PyTorch model we’ll first need to find the cached model path, which was downloaded during ‘Preflight’:

export MODEL_PATH=$(python -c "from huggingface_hub import snapshot_download; print(snapshot_download('Helsinki-NLP/opus-mt-en-fr'))")

We can then convert a model from PyTorch to ONNX:

python Marian-ONNX-Converter/convert.py $MODEL_PATH --no-quantize

We have to explicitly choose no quantization at this step, otherwise Marian-ONNX-Converter will attempt to quantize the encoder and decoder models in ways that are not supported by RKNN. It would also be premature to quantize the model before establishing baseline performance metrics. This will be the topic of part 3.

Note: I’m currently maintaining a fork of Marian-ONNX-Converter, which includes various small fixes that make it more appropriate for this project. This can be found here.

Model Files

Once we’ve converted the model to ONNX format, the output can be found in a sub-directory under outs/. In my case this was called outs/dd7f6540a7a48a7f4db59e5c0b9c42c8eea67f18, which was derived from the original model ID.

The contents of that directory should look something like this:

% ls -gho
total 1070376
-rw-r--r--  1    1.5K 13 Feb 21:59 config.json
-rw-r--r--  1    214M 13 Feb 21:59 decoder.onnx
-rw-r--r--  1    190M 13 Feb 21:59 encoder.onnx
-rw-r--r--  1    234K 13 Feb 21:59 lm_bias.bin
-rw-r--r--  1    116M 13 Feb 21:59 lm_weight.bin
-rw-r--r--  1    760K 13 Feb 21:59 source.spm
-rw-r--r--  1     74B 13 Feb 21:59 special_tokens_map.json
-rw-r--r--  1    784K 13 Feb 21:59 target.spm
-rw-r--r--  1    816B 13 Feb 21:59 tokenizer_config.json
-rw-r--r--  1    1.4M 13 Feb 21:59 vocab.json

As you might have guessed, the decoder.onnx and encoder.onnx files are ONNX models. These contain the model graphs and weights for the decoder and encoder, respectively. This is where attention happens!

The lm_weight.bin and lm_bias.bin files contain the weights and biases for the LM head. The raw output of a transformer is generally not in a form that is ready for human consumption. The LM head is a feed-forward neural network that converts transformer output into tokens that have meaning for humans.

The source.spm and target.spm files contain SentencePiece models for tokenization, and the vocab.json file ensures that tokens that occur in both the source and target languages are used efficiently. All other files in this directory are configuration files.

Converting to RKNN

Conversion from ONNX to RKNN format is relatively straight-forward, thanks to the RKNN Toolkit. We can simply load an ONNX model, then export that same model to RKNN format using the rknn.api module:

from rknn.api import RKNN

rknn = RKNN()
rknn.load_onnx(model='model.onnx')
rknn.build(do_quantization=False)
rknn.export_rknn('model.rknn')

We also want to convert the LM weights and biases from PyTorch format to raw floating point values:

def convert_weights(input_path, output_path):
    """Convert Torch LM head tensors to raw float32 for runtime matmul + bias.

    These can be loaded back into tensors in Python, or natively in C++.
    """
    tensor = torch.load(input_path, weights_only=True)
    weights = tensor.detach().numpy().astype(np.float32)
    weights.tofile(output_path)
    return weights

This is not strictly necessary, because our Python code can load the original lm_bias.bin and lm_weight.bin files directly, using PyTorch. However when we port the inference engine to C++, we will need raw floating point values. By converting now, we can be sure that our implementations behave consistently.

This conversion is handled by the rknn_convert.py script in my marian-rknn repo.

Quantization

In the example above, we’ve set do_quantization=False, which converts the model graph directly to RKNN format, with no quantization.

Quantization for encoder/decoder architectures is somewhat involved. For the encoder, we need to build a dataset containing a pair of NumPy files (the input IDs and attention mask) for each example. We then write the file names to a text file, one example per line.

Here’s what a dataset text file might look like:

calibration/input_ids_1.npy calibration/attention_mask_1.npy
calibration/input_ids_2.npy calibration/attention_mask_2.npy
calibration/input_ids_3.npy calibration/attention_mask_3.npy
calibration/input_ids_4.npy calibration/attention_mask_4.npy
calibration/input_ids_5.npy calibration/attention_mask_5.npy
...

The path to this file is then used as an argument for rknn.build():

rknn.build(do_quantization=True, dataset=calibration_path)

Effective quantization requires careful profiling and evaluation of a model, to ensure that performance does not suffer unnecessarily. I mention this mostly for completeness: for this implementation, we will leave quantization disabled and prioritize a working end-to-end RKNN pipeline.

Dynamic Ops

The next challenge we have is dynamic inputs. The Rockchip NPU has limited support for dynamic graph operations, which are required to handle inputs of arbitrary length. To avoid off-loading inference onto the CPU, we need to alter the model to use static input shapes.

For MarianMT we define these as:

ENCODER_INPUTS = ['input_ids', 'attention_mask']
ENCODER_INPUT_SIZE_LIST = [[1, 32], [1, 32]]

DECODER_INPUTS = ['input_ids', 'attention_mask', 'encoder_hidden_states']
DECODER_INPUT_SIZE_LIST = [[1, 32], [1, 32], [1, 32, 512]]

The array [1, 32] specifies an input length of 32 tokens for input_ids and attention_mask. The array [1, 32, 512] specifies an internal state with a shape of 32x512. The value 512 is not arbitrary - it is the size of the embedding used, which corresponds to d_model in the model’s config.json file.

The constants above are used when exporting the encoder and decoder models:

# Load encoder model
rknn = RKNN()
rknn.config(target_platform=platform,
            dynamic_input=[ENCODER_INPUT_SIZE_LIST])
rknn.load_onnx(model=encoder_path,
               inputs=ENCODER_INPUTS,
               input_size_list=ENCODER_INPUT_SIZE_LIST)

# Export encoder (trimmed for brevity)
rknn.build(...)
rknn.export_rknn(...)
rknn.release()

# Load decoder model
rknn = RKNN()
rknn.config(target_platform=platform,
            dynamic_input=[DECODER_INPUT_SIZE_LIST])
rknn.load_onnx(model=decoder_path,
               inputs=DECODER_INPUTS,
               input_size_list=DECODER_INPUT_SIZE_LIST)

# Export decoder (trimmed for brevity)
rknn.build(...)
rknn.export_rknn(...)
rknn.release()

Inference Engine

To keep things simple, our inference engine will be written in Python. Although there are potential benefits to writing an inference engine in C++, starting with Python has three main advantages:

It will be faster to iterate. We do not need to worry about memory management, linking, and other C++ footguns.
Interactions with SentencePiece will be closer to the original Hugging Face implementation, which will be helpful for comparison.
Finally, we can take advantage of simulation mode in the RKNN Toolkit, to run the model on your development machine. This is much faster than transferring .rknn files to device for testing.

This section will be divided into five steps, introducing code snippets for each part of the translation process:

Initialisation
Tokenization
Encoder
Decoder Loop
Detokenization

Let’s dive in!

Initialisation

During initialisation, we have to load each of the components of the model. This begins with two SentencePiece tokenisation models - one for the source language, and another for the target language:

# source language
spm_src = spm.SentencePieceProcessor(model_file=f"{input_path}/source.spm")

# target language
spm_tgt = spm.SentencePieceProcessor(model_file=f"{input_path}/target.spm")

We also need to load a vocab file, which maps token substrings to integers:

# load vocab from json file
vocab, vocab_inv = load_vocab(f"{input_path}/vocab.json", vocab_size)

Finally, we need to load the encoder and decoder RKNN models using RKNNLite:

# load encoder
rknn_enc = RKNNLite()
rknn_enc.load_rknn(f"{args.model_path}/encoder.rknn")

# load decoder
rknn_dec = RKNNLite()
rknn_dec.load_rknn(f"{args.model_path}/decoder.rknn")

Tokenization

You may wonder why we need two tokenizers and a vocabulary file. Surely a tokeniser would be sufficient for converting an input string into a series of integer token IDs…

There is a good reason for this. Marian allows the source and target languages to use the same token ID space. Say you have a token that appears in both the source and target languages (e.g. th), it can be more efficient to train the model using the same token ID for both languages.

The problem is that the two SentencePiece models are trained independently, leading common tokens to have different token IDs. The purpose of the vocab file is to correct this, by mapping token substrings to a common token ID space. This is best illustrated with a code snippet:

# break input string into token substrings
pieces = spm_src.encode(line, out_type=str)

# map token substrings to token IDs
tokens = [vocab.get(piece, unk_token_id) for piece in pieces]

These token IDs can then be used as input to the encoder.

Encoder

The encoder begins by building a fixed-length input_ids tensor. This is padded or truncated to the appropriate input length. The “end of stream” and “pad” tokens are part of the model configuration. We also need to build an attention_mask, in which valid token positions are marked as 1 and padded positions as 0. This gives the model visibility into which positions are meaningful and which are merely padding.

# prepare encoder inputs and attention mask
encoder_input_ids = prepare_encoder_inputs(tokens, enc_len, pad_token_id, eos_token_id)
attention_mask = build_attention_mask(encoder_input_ids, eos_token_id)

To invoke the encoder, we call rknn_enc.inference(...) with both arrays in the order expected by the model:

# invoke the encoder
encoder_outputs = rknn_enc.inference(inputs=[encoder_input_ids, attention_mask])
encoder_hidden_states = encoder_outputs[0]

The return value is a tensor (named encoder_hidden_states) which contains an embedded representation of each source token position. This will be used by the decoder to extract meaning from the source text.

Decoder Loop

Unlike the encoder, which runs just once, the decoder runs in an autoregressive loop. This means that each token generated depends on the output from previous iterations of the decoder.

The decoder requires three input tensors:

Decoder input IDs (not the same as the input IDs passed to the encoder)
The same attention mask prepared earlier
The encoder_hidden_states produced by the encoder

The first input is essentially the decoder’s “working memory”. On each iteration of the decoder, the model may attend to tokens that were previously generated. This helps inform which token is chosen next.

The third input is simply the output of the encoder.

We have a choice of two different decoding strategies, which are discussed further below:

if beam_search is None:
    # simpler, but input may be of slightly lower quality
    output_tokens = greedy_decode(
        rknn_dec,
        attention_mask,
        encoder_hidden_states,
        ...
    )
else:
    # more complex, but better output quality
    beam_width, beam_depth = beam_search
    output_tokens = beam_search_decode(
        rknn_dec,
        attention_mask,
        encoder_hidden_states,
        ...
    )

Greedy decoding is simpler, so we’ll focus on that for now. This snippet of code shows the entire decoder loop for greedy decoding:

def greedy_decode(
    rknn_dec,
    attention_mask,
    encoder_hidden_states,
    weight,
    bias,
    decoder_start_token_id,
    pad_token_id,
    eos_token_id,
    dec_len
):
    # prepare autoregressive decoder input; start by fully padding the tensor
    decoder_input_ids = np.full((1, dec_len), pad_token_id, dtype=np.int32)
    decoder_input_ids[0, 0] = decoder_start_token_id

    # this will hold our final decoder output
    output_tokens = []

    # main loop
    for step in range(dec_len - 1):
        # invoke decoder, passing in the three inputs
        outputs = rknn_dec.inference(inputs=[decoder_input_ids, attention_mask, encoder_hidden_states])

        # apply LM head, which converts the output into logits
        decoder_output = outputs[0]
        hidden = decoder_output[0, step, :].astype(np.float32)
        logits = hidden @ weight.T + bias

        # greedily decode next token by choosing the candidate
        # with the highest logit
        next_token = int(np.argmax(logits))
        output_tokens.append(next_token)

        # check for end of output
        if next_token == eos_token_id:
            break

        # otherwise, append token to decoder input for next iteration
        if step + 1 < dec_len:
            decoder_input_ids[0, step + 1] = next_token

    return output_tokens

The output of the decoder is a sequence of token IDs. The final step is to turn those token IDs into human-readable text!

Detokenization

Detokenization begins by inverting the vocab mapping applied earlier - converting token IDs into substrings. Those substrings are passed to the target tokenizer to be converted into human-readable text:

decoded_pieces = []
for token_id in output_tokens:
    if token_id in (eos_token_id, pad_token_id) or token_id <= 0:
        break
    piece = vocab_inv.get(token_id)
    if piece is not None:
        decoded_pieces.append(piece)

return spm_tgt.decode(decoded_pieces)

Decoding Strategies

In the decoder loop section above, we mentioned having two decoding strategies to choose from. These are Greedy Decoding and Beam Search. This choice is all about trade-offs.

Greedy decoding is fast and deterministic, picking the highest probability token at each step. Beam search explores multiple hypotheses in parallel, which often improves fluency or adequacy. However it also increases compute, memory pressure, and book-keeping costs.

Greedy Decoding

As tokens are generated, greedy decoding uses np.argmax(logits) to choose the highest probability token to append to the output.

An interesting observation we can make here is that logits are not the same as probabilities. Normally, we would use a Softmax function to convert the logits to probabilities:

# calculate softmax
logits -= np.max(logits)
exp_logits = np.exp(logits)
probs = exp_logits / np.sum(exp_logits)

In the case of greedy decoding, we can rely on the fact that argmax(logits) will identify the same token as argmax(probs).

Although efficient, the downside of greedy search is that the decoder cannot backtrack on tokens that lead to poorer translations.

Beam Search

Beam search addresses this by managing multiple candidate sequences at each step. The number of concurrent sequences (or beams) is known as the beam width.

To decide which beams should be extended, and which should be culled, the algorithm maintains a cumulative probability. After expanding each candidate beam with possible next tokens, only the top-ranked hypotheses are preserved.

Image Credit: Chapter 10.8 - Beam Search from Dive Into Deep Learning.
This diagram shows a beam size of 2 and maximum sequence length of 3. Greedy search can be treated as a special case of Beam Search where the beam size is set to 1.

Beam search reduces the risk of early mistakes dominating the final output. But it comes at the cost of compute, memory pressure and additional book-keeping. This extra work is primarily in the decoder and LM-head, which must be executed once per beam on each iteration.

Benchmarks

The final step in this process is benchmarking. How many sentences can the RK3588 translate per minute? Or more precisely, how many tokens?

Sample Text

To prepare for the benchmark, we need a file containing input sentences that are representative of what our system will encounter under real world conditions. The Random English Sentences dataset from Kaggle is a good place to start! Here’s what the first 10 lines look like:

I didn't know what to wear to my best friend's funeral.
There used to be a big cherry tree behind my house.
Sometimes, books and movies made her feel things she couldn't explain.
My bandaid wasn't sticky anymore so it fell off on the way to school.
He only had one drink but he's completely wasted.
In Hong Kong, everyone lives stacked together.
Her parents had been separated for almost three years, but she never expected them to actually go through with the divorce.
These are the things that I want from a relationship.
When one considers the issue from a different perspective, the answer seems to change.
She twisted her ankle and couldn't go to soccer anymore.

Benchmark Script

The marian-rknn repo includes a benchmark script (benchmark.py) which repeatedly translates input for a fixed period of time. Once complete, it prints throughput and token count statistics.

To run the benchmark, we have to provide a model path, a dataset, and a maximum runtime (in seconds). The benchmark will iterate over the dataset until the time budget is reached, then report aggregate throughput and per-stage timings:

python scripts/benchmark.py \
  outs/dd7f6540a7a48a7f4db59e5c0b9c42c8eea67f18 \
  datasets/en-phrases.txt \
  120 > results-greedy.txt

We can run the same test using Beam Search:

python scripts/benchmark.py \
  --beam-search \
  --beam-width 3 \
  outs/dd7f6540a7a48a7f4db59e5c0b9c42c8eea67f18 \
  datasets/en-phrases.txt \
  120 > results-beam-search.txt

Results

Here are the results!

As a baseline, greedy decoding with the datasets/en-phrases.txt dataset allowed us to produce 2766 output tokens. In total, we processed 207 sentences @ 1.705 sentences per second:

Benchmark complete
Elapsed: 121.442 s
Sentences: 207
Sentences/sec: 1.705
Total time: 121295.378 ms
Encoder time: 2492.285 ms
Decoder time: 91355.883 ms
LM head time: 26394.237 ms
Avg total time per sentence: 585.968 ms
Avg encoder time per sentence: 12.040 ms
Avg decoder time per sentence: 441.333 ms
Avg LM head time per sentence: 127.508 ms
Input tokens: 2379
Output tokens: 2766
Decoder iterations: 2766
Input tokens/sec: 19.590
Output tokens/sec: 22.776

When we run the same test with beam search enabled, and a beam width of 3, we see how much of an impact it can have…

Benchmark complete
Elapsed: 121.982 s
Sentences: 52
Sentences/sec: 0.426
Total time: 121851.846 ms
Encoder time: 631.315 ms
Decoder time: 72169.512 ms
LM head time: 20800.915 ms
Avg total time per sentence: 2343.305 ms
Avg encoder time per sentence: 12.141 ms
Avg decoder time per sentence: 1387.875 ms
Avg LM head time per sentence: 400.018 ms
Input tokens: 663
Output tokens: 729
Decoder iterations: 2177
Input tokens/sec: 5.435
Output tokens/sec: 5.976

Average decoder time per sentence goes up considerably, and the number of output tokens produced drops from 2766 to 729. Overall, we processed 52 sentences @ 0.426 sentences per second.

It’s worth noting that this is without any attempt at further optimisation.

Wrap Up

At this point, we have a complete path from a pretrained MarianMT model to inference on the Rockchip NPU: exporting to ONNX, converting the encoder and decoder to RKNN, running tokenization and detokenization, and comparing greedy decoding with beam search.

There is still plenty of room for refinement. Quantization can be used to reduce memory usage and compute overhead, while further optimisation of the decoder loop and LM head may improve throughput. If you want to follow future progress, keep an eye on the marian-rknn repo we’ve used throughout this post, where I will continue to document quantisation and optimisation efforts.

References

Understanding SentencePiece
Using dynamic input shapes on RKNN/RK3588 by Martin Chang
Marian-ONNX by Kevin Costa
Unlocking On-Device AI by John OSullivan
RKNN API Demo by Rockchip
Quantize ONNX Models

Machine Translation at the Edge: Part 1

Mon, 05 Jan 2026 00:00:00 +0000

This is my third post exploring Edge AI, with a focus on the Rockchip RK3588 NPU. This post dives into Machine Translation, a Natural Language Processing task that has traditionally been associated with Cloud AI services. With advances in Edge AI hardware, model architectures and distillation techniques, it is now feasible to perform Neural Machine Translation directly on an embedded device.

The goal of this post (and the next) is to build a Neural Machine Translation pipeline for the Rockchip NPU. We’ll walk through model selection, quantization, and deployment using the RKNN Toolkit. Python and C++ source code will be available on GitHub. The initial implementation is based on the Khadas Edge2 development board, but should be easy to port to other RKNN-based devices.

If you haven’t already, I recommend reading the first and second posts, as they cover helpful background material.

Machine Translation
RKNN Platform
RKNN Model Zoo
MarianNMT vs MarianMT
Test Flight
Next Steps
References

Machine Translation

Machine Translation is a technique used to automatically, or mechanically, translate text (or speech) from one language to another. We generally refer to the two languages involved as the ‘source’ and ‘target’ languages. At a high level, the machine reads text in the source language, and generates a corresponding translation in the target language.

Machine Translation has been an active area of research since the early 1950s. In fact, one of the first conferences on Mechanical Translation (as it was known at the time) was hosted at MIT in 1952. With such a rich history, there are many translation models to choose from. We’ll briefly review some earlier architectures here.

Statistical Models

Human language is highly redundant. After seeing the first few letters of a word we can often predict the following letters with fairly high accuracy. The same is true for words in a sentence, and even phrases. This motivates the treatment of language as probabilistic in nature.

In Statistical Machine Translation (SMT), a model learns a probabilistic mapping from a source language to a target language. A classic example of this is Phrase-Based Machine Translation (PBMT). A sentence is first segmented into phrases, and each phrase is mapped to corresponding phrases in the target language:

Image credit: Adapted from the original image found here.
This was included in the ACL 2005 Workshop on Building and Using Parallel Texts.

Once candidate word/phrase matches have been selected, the system searches for the best overall translation by ranking and re-ordering the translations, before combining them to produce a fluent result.

Aside: Moses

If you start digging into the literature around statistical models, you are likely to stumble upon Moses, an open-source SMT system. Although Moses is no longer common in front-line translation systems, it is still relevant for ‘low-resource’ language research (i.e. translating languages with limited training data) and as a teaching tool.

You may also discover Sacremoses, a Python port of Moses’ text preprocessing scripts. This is often used in older model pipelines, or as an optional dependency. SentencePiece, BPE and Unigram tokenizers are more common these days.

The primary issue with SMT systems is that they depend on a pipeline of hand-crafted components. Because translation occurs across multiple stages (e.g. segmentation, alignment, phrase tables, and reordering) errors in early stages can easily propagate to the final translation. The result ends up being brittle and expensive to maintain.

Encoder-Decoder Models

Encoder-Decoder models replace the individual components of SMT and PBMT with a single end-to-end model that can learn directly from a large collection of source-target language pairs. This works remarkably well on high-resource languages (i.e. those with plenty of high-quality training data).

Another reason that Encoder-Decoder architectures have been so successful is that Machine Translation is fundamentally a sequence-to-sequence problem - you take a variable-length source sentence and produce a target sentence. The source and target sentences will often be different lengths, since an expression in one language may look very different in another.

Even within the sub-category of Encoder-Decoder models, we have a few variations to consider… the first are Recurrent Neural Networks.

RNNs and LSTMs

Recurrent Neural Networks (RNNs) are designed for processing sequential data, where the order of elements is important. This is clearly the case for language, where the meaning of a word depends on the presence and order of words (or more generally, tokens) around it.

Unlike feed-forward networks (FFNs), which process inputs independently, RNNs utilize recurrent connections. These recurrent connections carry ‘hidden state’ from one iteration to the next - allowing the network to progressively capture the meaning in a sentence. This is well suited to translation, where the source sentence is first processed by an Encoder RNN to extract the meaning, then passed to a Decoder RNN to generate output tokens:

Some information is passed on from each iteration of an RNN to the next - this is the hidden state of the RNN. In this architecture, the source sentence is processed by the encoder RNN, until the final token is seen (</code>, or ‘End of Sequence’). The final hidden state is then used to initialise the decoder RNN, which generates tokens until </code>.</small>

RNNs do suffer from some limitations, such as inherently sequential execution, and difficulty with long-range dependencies. They also have a fixed-sized hidden state vector, which limits representational capacity. Since an entire sentence must be squeezed into a fixed-size vector, RNN-based models may struggle with longer sentences.

Transformers

The most recent Machine Translation models tend to use some variation of the Transformer Architecture, which we covered in the previous post on Tiny Language Models.

Image Credit: Wikipedia.

Unlike RNNs, Transformers do not rely on recurrent hidden state. They instead use Self-attention, which allows each token to incorporate information from all other tokens in the sequence. This is implemented in such a way that all tokens can be processed in parallel, which addresses one of the major bottlenecks of RNNs.

After the encoder has processed a sentence, the decoder typically generates one token at a time, conditioned on the previously generated tokens.

Together, these characteristics make the Transformer architecture particularly well suited to Machine Translation, and this is the kind of architecture we’ll be discussing throughout the rest of this exercise.

RKNN Platform

Let’s move on to our target platform - this includes our choice of our hardware and framework support.

Why use the Edge2?

The target for this exercise is the Khadas Edge2. This compact hobbyist board is based on the Rockchip RK3588 SoC, an 8-core 64-bit ARM processor with a dedicated neural processing unit (or NPU).

The RK3588 SoC uses a Unified Memory Architecture (UMA). In practice, this means the CPU and NPU operate on memory that is allocated from the same pool of system memory. This is different to discrete desktop GPUs, where the CPU and GPU each have their own pool of memory and often results in large data transfers over PCIe.

Platform Limitations

While UMA simplifies communication between the CPU and NPU, we need to be mindful of allocation constraints (e.g. contiguous blocks of memory, required for DMA) and monitor overall memory pressure. Large models can easily lead to out-of-memory errors and allocation failures.

Bandwidth is also shared, so concurrent CPU and NPU workloads compete for DRAM throughput even if memory capacity is sufficient. Our preference is to find smaller models that are specialised to the task at hand, rather than more generalised large language models.

RKNN Toolkit

To take advantage of hardware acceleration on the Rockchip NPU, we need to interface with the RKNPU driver via the RKNN API. Rockchip’s RKNN Toolkit is an SDK for the Rockchip NPU that supports model conversion in Python (on PC), and inference/evaluation on Rockchip devices.

The toolkit supports conversion, quantization and fine-tuning models on PC. It also includes runtime and C/C++ headers that we need to interface with the NPU in C++:

The RKNN Toolkit includes components that run on-device (for inference), and others that are intended to be run on a PC (for model construction and conversion).

A typical RKNN workflow begins with a pretrained model in PyTorch, TensorFlow or ONNX format. RKNN Toolkit includes Python APIs to load models from other formats, and export them in RKNN format:

from rknn.api import RKNN

rknn = RKNN()
rknn.load_onnx(model='model.onnx')
rknn.build(do_quantization=True, dataset='dataset.txt')
rknn.export_rknn('model.rknn')

Once in RKNN format, a model can be loaded on-device using C++ or Python code.

Model Support

Although there are RKNN implementations for some popular model architectures, most pretrained models will be published in other formats. Common formats in the wild include PyTorch and TensorFlow. Another common format is ONNX, short for Open Neural Network Exchange. This has gained popularity as a model interchange format.

The RKNN API makes it easy to load and export models, but whether this will be sufficient depends heavily on the model architecture. In some cases, it may be impossible to run a model on the Rockchip NPU due to certain layers or operations not being supported. For example, models that rely on dynamic graph operations are not well supported.

When converting models, it can be helpful to visualise the model graph, in order to diagnose problems. Netron is a convenient and easy-to-use viewer for neural networks:

Image credit: Screenshot taken from the Netron GitHub repo.

RKNN Model Zoo

When implementing a machine learning model, it is generally a bad idea to start from scratch. Let’s look at some of the available examples.

The RKNN Model Zoo provides a range of examples for the Rockchip platform, primarily computer vision models. But there is one we can use for Machine Translation, based on a model architecture known as the Lite Transformer.

Lite Transformer

The Lite Transformer is a variant of the Transformer model that makes certain trade-offs to improve performance on resource constrained devices. The authors recognised that computation in a standard Transformer is dominated by the feed-forward network (FFN). This is despite the fact that the FFN comes after Attention, and so does not mix context across tokens.

A straightforward way to reduce compute demands is to shrink the embedding size directly - this reduces capacity across the entire model. Alternatively, we can reduce the attention feature dimension. While bottlenecking attention can save compute, the FFN still dominates compute, and the savings can be limited relative to the quality hit.

The authors introduce a clever solution to this problem, called Long-Short Range Attention.

Long-Short Range Attention

Long-Short Range Attention (LSRA) splits the representation into two parallel branches - a convolutional neural network, specialised for learning short-range dependencies between tokens, and an attention block that can focus on long-range dependencies.

Image Credit: Figure 3 from Lite Transformer with Long-Short Range Attention by Wu et al. This shows how conventional attention predominantly captures local (short-range) dependencies (b). In LSRA, the presence of the convolutional network for local context allows the attention mechanism to better capture long-range dependencies (c).

The Lite Transformer authors use a threshold of 500M multiply-add operations to measure performance of the model in a mobile setting. By configuring various models to operate within this threshold, they showed that LSRA often outperforms pure transformer models in resource-constrained environments.

Other Choices

While models such as Lite Transformer have great potential for reducing computational cost on embedded devices, there are some trade-offs we need to consider. The two we’re most concerned about are availability of pretrained models, and additional challenges in fine-tuning.

For this post, I’ve shortlisted several alternatives that directly address the needs for model availability and fine-tuning:

Google’s T5 is a versatile text-to-text transformer, capable of a variety of NLP tasks, including translation and summarization. T5 uses a generalised model that needs to be prefixed with the specific task (e.g. "Translate English to French: I am a fish."). Possibly too big for our embedded use-case, especially if we only need to support one language pair.
MarianNMT is a pure C++ toolkit specialised for training Neural Machine Translation (NMT) models. It is often faster in inference, and capable of producing higher quality models for specific language pairs. Marian has been used in production systems by Microsoft and Mozilla (via their own fork). Strong contender!
OpenNMT is similar to MarianNMT in that it provides an ecosystem for NMT and neural sequence learning. The OpenNMT project also maintains CTranslate2, an efficient inference engine for transformer models. Also a strong contender!

While CTranslate2 includes support for AArch64 (ARM 64-bit) CPUs, our ideal target is the Rockchip NPU. Both OpenNMT and MarianNMT will require a custom inference implementation and model conversion process. Due to the availability of model checkpoints for Marian, that’s where we’ll focus our attention going forward.

MarianNMT vs MarianMT

As mentioned above, MarianNMT is a toolkit specialised for training NMT models. Somewhat confusingly MarianMT (without the N) refers to a collection of Transformer-based encoder–decoder models, published on Hugging Face.

All MarianMT variants follow the Transformer architecture, with 6 layers in both encoder and decoder. They are all natively ONNX exportable, which is good for us, because it eases the transition to RKNN format.

For training, MarianMT relies on the MarianNMT framework, which is written in C++. Marian can run on the CPU, however GPU acceleration is available via CUDA/cuDNN when enabled.

The MarianNMT project includes an efficient inference engine, also written in C++. This can be compiled to run on the CPU, or on an NVIDIA GPU using CUDA. We can also use an implementation based on Hugging Face Transformers.

Build from Source (Optional)

Getting started with MarianNMT is relatively straightforward… assuming you have a Linux system with an NVIDIA GPU, and the necessary CUDA Toolkit / cuDNN libraries installed.

To begin with, you can clone the MarianNMT source code from GitHub:

git clone https://github.com/marian-nmt/marian.git
cd marian

Then we can build it using CMake:

mkdir build && cd build
cmake ..
make

This builds the entire toolkit, which includes tools for training and the inference engine.

The MarianNMT build depends on CUDA by default, so you won’t have much luck if you try to build this on macOS or another platform without CUDA support. To try compiling without CUDA support:

cmake .. -DUSE_SENTENCEPIECE=on -DCOMPILE_CPU=on -DCOMPILE_CUDA=off
make

This will be somewhat slower, but good enough for a simple test run.

Note: This may not work on newer Macs based on Apple Silicon. Don’t worry if this means you cannot run the examples yourself. I’ve included example output below, so you can follow along!

NPZ Format

Unfortunately, Hugging Face checkpoints aren’t directly usable by the Marian C++ decoder. The MarianNMT inference engine expects an .npz file…

The Marian .npz format is essentially a ZIP archive containing one or more .npy (NumPy) arrays, which store model weights in a format appropriate for Marian NMT models. It is based on NumPy’s .npz format and is specifically tailored for fast loading and execution in the C++ Marian decoder.

We can see how this works with a simple Python script:

import numpy as np

# Save arrays to an npz file
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
np.savez('arrays.npz', array1=a, array2=b)

# Load arrays from an npz file
data = np.load('arrays.npz')
print(data['array1'])  # Output: [1 2 3]
print(data['array2'])  # Output: [4 5 6]

One way we can address this is to convert a pretrained model from another format (e.g. PyTorch) to .npz format. However, this can be error-prone, and it’s an unpleasant way to begin experimenting with a new architecture.

Pretrained Models

Alternatively, we can find pretrained models that already have .npz files available. There are a range of models available, trained on various language pairs and datasets. Common datasets for Machine Translation include:

WMT - WMT (short for ‘Workshop on Machine Translation’) is a shared-task dataset created annually for Workshop on Machine Translation competitions. Includes major language pairs and is trained on news, formal text, and curated sources.
Tatoeba - A community driven dataset consisting of example sentences and translations, all provided by volunteers. This includes hundreds of languages, including many ‘low resource’ languages.
OPUS - OPUS (Open Parallel Corpus) is an umbrella project hosting hundreds of parallel corpora from many domains, scraped from the web or public sources. Great for translation, as it contains billions of sentence pairs across 400+ languages.

Test Flight

For the rest of this post, we’ll focus on an English <-> French model that was pretrained on the OPUS dataset.

Getting the Model

The specific model we want is en-fr from the OPUS-MT-train repo on GitHub. The model is a few years old now, having been trained in 2020. However, it’s in NPZ format, ensuring that we can use it directly with MarianNMT.

The model’s README file provides a link to a download mirror. Let’s download and extract the model:

wget https://object.pouta.csc.fi/OPUS-MT-models/en-fr/opus-2020-02-26.zip
unzip -d opus-2020-02-26 opus-2020-02-26.zip

Note: The instructions below assume you’ve extracted the model to marian/build/opus-2020-02-26.

The zip file will be about ~265MB, and the extracted model will take about 288MB on disk. The file list when extracting may be a bit intimidating at first:

Archive:  opus-2020-02-26.zip
  inflating: opus-2020-02-26/README.md
  inflating: opus-2020-02-26/LICENSE
  inflating: opus-2020-02-26/opus.spm32k-spm32k.transformer-align.model1.npz.best-perplexity.npz
  inflating: opus-2020-02-26/opus.spm32k-spm32k.vocab.yml
  inflating: opus-2020-02-26/opus.spm32k-spm32k.transformer-align.valid1.log
  inflating: opus-2020-02-26/opus.spm32k-spm32k.transformer-align.train1.log
  inflating: opus-2020-02-26/source.spm
 extracting: opus-2020-02-26/source.tcmodel
  inflating: opus-2020-02-26/target.spm
  inflating: opus-2020-02-26/decoder.yml
  inflating: opus-2020-02-26/preprocess.sh
  inflating: opus-2020-02-26/postprocess.sh

There are several files that are worth pointing out. The first is opus.spm32k-spm32k.transformer-align.model1.npz.best-perplexity.npz, which contains Marian model weights in NPZ format. This is the core of the model. The string spm32k tells us that the specific tokenizer is SentencePiece (spm) with a vocab size of 32,000.

The next two files of interest are source.spm and target.spm, which are SentencePiece models. These models contain the information needed to segment an input sentence into tokens (i.e subword parts), and to concatenate generated tokens into the final output. This leads us to a more in-depth discussion on tokenization…

Tokenization

When translating a piece of text, that input must first be segmented into units of text called tokens. Tokens may be entire words, individual characters, or groups of characters (i.e. subwords):

Token granularity is an implementation detail that must be included in a model definition. MarianNMT often uses SentencePiece to tokenize text into subwords. SentencePiece itself supports the Unigram and Byte-Pair Encoding (BPE) algorithms. Which one is used is model-dependent.

To make this more concrete, we can use Marian’s spm_encode to tokenize an input sentence ("I am a spiky fish") using the source tokenizer:

$ echo "I am a spiky fish" | ./spm_encode --model opus-2020-02-26/source.spm
I ▁am ▁a ▁sp ik y ▁fish

We can see that the output is a series of tokens (or subword parts), but our model requires integer token IDs. The last file we’re interested in is opus.spm32k-spm32k.vocab.yml - this provides a ‘vocabulary’, mapping tokens to integers (and vice-versa).

Vocab File

The vocab YAML file is human readable, so we can take a peek inside… This is the first 20 lines of opus.spm32k-spm32k.vocab.yml:

$ head -n 20 opus-2020-02-26/opus.spm32k-spm32k.vocab.yml
</s>: 0
<unk>: 1
",": 2
.: 3
▁the: 4
▁de: 5
"'": 6
▁of: 7
▁la: 8
s: 9
▁and: 10
▁et: 11
▁to: 12
▁des: 13
▁l: 14
▁a: 15
▁les: 16
▁à: 17
▁in: 18
▁le: 19

The list begins with </s> and <unk>, which represent end-of-sequence and unknown strings, respectively. Unknown strings occur when the input contains character sequences that do not have a valid mapping.

Next, we see some common symbols , and ., followed by strings that look like parts of words. These are often prefixed with a whitespace marker ▁, which bears a striking resemblence to an underscore… The whitespace marker means that these are prefixes - the first part of a word that occurs after a space. This distinction between subwords that follow a space and those that don’t has been shown to improve model performance.

Full Pipeline

With all that covered, we can invoke the model using a Unix pipeline consisting of three executables: spm_encode, marian-decoder and spm_decode.

We can invoke the decoder using this rather complicated command:

echo "I am a fish" \
  | ./spm_encode \
      --model opus-2020-02-26/source.spm \
  | ./marian-decoder \
      --models \
        opus-2020-02-26/opus.spm32k-spm32k.transformer-align.model1.npz.best-perplexity.npz \
      --vocabs \
        opus-2020-02-26/opus.spm32k-spm32k.vocab.yml \
        opus-2020-02-26/opus.spm32k-spm32k.vocab.yml \
      --devices 0 \
      --beam-size 6 \
      --normalize 1.0 \
  | ./spm_decode \
      --model opus-2020-02-26/target.spm

Starting at the top, we use echo to output the string "I am a fish". This is piped to spm_encode so that SentencePiece can tokenize the input. Tokens are then piped to marian-decoder.

Note that when we invoke marian-decoder, the same vocab file is provided twice - this is not a typo. Marian models support asymmetric translation - i.e. using different vocabularies for the source language (encoder) and target languages (decoder). Therefore, the marian-decoder command expects paths for both source and target vocabulary files. As we saw above, this file contains a mix of English and French tokens, so we can use the same file for both.

Finally, the output of marian-decoder is piped to spm_decode to be converted into human-readable text!

[2026-01-04 19:53:04] [config] Loaded model has been created with Marian v1.8.2 2111c28 2019-10-16 08:36:48 -0700
...
[2026-01-04 19:53:05] Loaded model config
[2026-01-04 19:53:05] Loading scorer of type transformer as feature F0
[2026-01-04 19:53:05] [memory] Reserving 284 MB, device gpu0
[2026-01-04 19:53:05] [gpu] 16-bit TensorCores enabled for float32 matrix operations
[2026-01-04 19:53:05] Best translation 0 : ▁Je ▁suis ▁un ▁poisson
Je suis un poisson
[2026-01-04 19:53:05] Total time: 0.05016s wall

If we comb through the output, we find the string ‘Je suis un poisson’, which is a valid French translation of ‘I am a fish’!

Next Steps

That completes our introduction to machine translation and the MarianNMT toolkit. We’ve covered a lot of ground! However, everything we’ve done so far has been on PC, rather than our target device.

In part two, we turn our attention to the Rockchip NPU. Our goal will be to perform MarianMT inference on the Rockchip NPU. We will download a pretrained MarianMT model from Hugging Face, convert it to RKNN format, then run it directly on Rockchip NPU.

To make all of this work, we’ll need to consider quantization, input padding and layer compatibility. Thankfully, the RKNN Python API makes this much easier than it sounds. Our result will be a Marian implementation that runs directly on the Rockchip NPU!

References

Historical:

MIT Conference on Mechanical Translation - Mostly unpublished papers from the early days of Machine Translation.
Moses SMT - This is worth exploring, if only to appreciate the extensive research that went into constructing Machine Translation systems prior to the rise of Transformers.

NMT Systems:

Tokenization:

Byte-Pair Encoding - Hugging Face
Unigram - Hugging Face
Understanding SentencePiece - A nice introduction to SentencePiece.

Tiny Language Models

Sun, 07 Dec 2025 00:00:00 +0000

This post continues on the theme of Edge AI. We previously covered the basics of Edge AI and the practicalities of running YOLO, an Object Detection algorithm, on the Rockchip NPU. In this post, we’ll look at the history and scaling challenges of Large Language Models and what it means to train and deploy a tiny language model.

Although there is an example for the Khadas Edge2 at the end, the content here is mostly foundational in nature. In a later post, we will look at practical examples of running Machine Translation algorithms on the Rockchip NPU.

A Brief History of Language Models…
Large Language Models
Transformer Architecture
Tiny Language Models
Tokenization
RKLLM Toolkit
Next Steps…
- Machine Translation at the Edge
- Tokenization Deep Dive
References

A Brief History of Language Models…

The history of language models (or more generally, Natural Language Processing) features some clearly visible paradigm shifts. In the early days of computing, computer scientists and programmers were quick to explore the AI landscape. As early as the 1960s, chatbots such as of Joseph Weizenbaum’s ELIZA led people to believe that the era of machine intelligence was upon us.

Fast forward to the 1980s and 90s, and by this time N-gram models had taken hold. These probabilistic models aimed to predict words and symbols that would occur in a text, using the surrounding context to condition the probability. Although the results were compelling, they had properties that made them very difficult to scale.

In the early 2000s, neural networks were successfully applied to language modeling. The immense representational power of neural networks enabled the development of word embeddings, like those seen in Word2Vec. Although vanilla neural networks quickly reached their limits, research in the 2010s gave us Recurrent Neural Networks and LSTMs, which allowed neurons to preserve state from one word/symbol to the next.

Finally, the invention of Transformers in 2017 ushered in the era of Large Language Models. You can see from the timeline below that we’re currently in a period of rapid growth, marked not just by leaps in computational capacity, but also paradigm shifts:

We’ll dive into each of these in a bit more detail.

N-gram Models

N-gram models are purely probabilistic language models, trained on a body of text (commonly known in NLP as a “corpus”) to capture the frequency and context in which different symbols occur. Symbols can be individual characters, words, or even sub-words (e.g. prefixes and suffixes).

N-gram models are characterised by the length N of the input sequence used to condition the probability of a symbol that follows in the sequence. Assuming that entire words are used as symbols:

Using Bayes’ rule (or Bayes’ theorem), we can calculate the probability of a given input sequence as a joint probability:

\[\begin{aligned} &P(\text{The}\text{ }{sky}\text{ }{is}\text{ }{blue}) \\ = &P({The}, {sky}, {is}, {blue}) \\ = &P({The}) * P({sky}\text{ }|\text{ }{The}) * P({is}\text{ }|\text{ }{The}, {sky}) * P({blue}\text{ }|\text{ }{The}, {sky}, {is}) \end{aligned}\]

N-gram models suffer from two critical limitations. First—they have limited expressive capabilities. Specifically, the probability of a particular symbol occurring may be dependent on context that falls outside the conditioning sequence. We can attempt to address this by increasing N to capture more context.

This leads to the second issue - as the value of N increases, the data and storage requirements also grow exponentially. And if the dataset is not sufficiently large and diverse, the model risks overfitting.

Neural Language Models

The introduction of Neural Language Models in 2003 also gave us the concept of Word Embeddings. A Word Embedding is a vector representation of a word, where each component in the vector corresponds to some concept captured by the model.

Let’s make this more concrete. Say you have a vocabulary of five words, and each word is represented using its one-hot encoding:

Cat    -> [1, 0, 0, 0, 0]
Kitten -> [0, 1, 0, 0, 0]
Dog    -> [0, 0, 1, 0, 0]
Puppy  -> [0, 0, 0, 1, 0]
Bird   -> [0, 0, 0, 0, 1]

If we used these vectors as inputs for a neural network, the input layer will require five neurons: one for each component of the input vector. Unfortunately, as the vocab grows in size, say to 40,000 words, the input layer of our neural network will also grow to be 40,000 neurons wide.

Besides scaling issues, one-hot encoding comes with two other disadvantages:

Sparsity - Most values in each vector are zeros. This is wasteful of storage, and means that the first layer of the network has relatively little information to work with.
No Relationships Between Words - The disjoint nature of one-hot vectors means that we cannot derive any meaningful relationships between words.

To highlight the second point, consider the first two examples: Cat and Kitten. A Kitten is a kind of Cat, so it doesn’t make sense for those to be completely disjoint in their representations. Word Embeddings address this by learning a dense representation (an embedding) for each word in the dataset. The embedding reduces the dimensionality of the input data, allowing the rest of the model to operate on vectors with a lower dimensionality.

Image Credit: Figure 1 from A Neural Probabilistic Language Model by Bengio et al (2003).
The major contribution of Bengio et al’s paper was train a neural language model end-to-end.

Word2Vec

A critical issue with Neural Language Models is that training at scale was expensive and suffered from tight coupling between network architecture and word representations. In 2013, Word2Vec showed that it was possible to learn meaningful relationships between words by training a shallow network on a large corpus. The resulting word embeddings could then be re-used across different architectures.

Image Credit: The Illustrated Word2Vec by Jay Alammar.
He goes into great detail about how Word2Vec can be used in a variety of applications.

Recurrent Neural Networks

Recurrent Neural Networks (RNNs) made it possible to preserve some state from one token to the next. This meant that generated tokens would take more context into account - for example, choosing the generated token based on the meaning implied by the most recent words in the input sequence.

Image Credit: A Brief Overview of Recurrent Neural Networks from Analytics Vidhya.

A key problem with RNNs is vanishing gradients. Long Short-Term Memory (LSTM) networks build upon RNNs by introducing a new internal structure…

Long-Short Term Memory Networks

Although the idea for Long Short-Term Memory (LSTM) networks was first described in 1997 by Hochreiter et al, it wasn’t until the 2010s that this found wider adoption in Language Models. LSTMs attempt to solve the vanishing gradient problem by ensuring that errors flow back through the network, while introducing gating mechanisms that allow the error (or gradient) to be truncated when it is beneficial.

However, this wasn’t quite enough to get us to where we are now… RNNs and LSTMs are inherently serial, which causes severe bottlenecks. Vanishing gradients are a problem, even in LSTM networks. This makes it increasingly difficult for models to learn connections between tokens over longer time spans.

Transformers

The NLP landscape changed dramatically in 2017, when the seminal paper Attention Is All You Need introduced the Transformer architecture. The Transformer architecture replaces recurrence with a mechanism called “Attention.” This made it possible for language models to capture long-range dependencies.

You’re almost certainly aware of Transformers (although you may not recognise them by name). They are the “T” in GPT (short for Generative Pre-trained Transformers) and have led to many state-of-the-art results in Natural Language Processing and Machine Learning tasks.

Large Language Models

Current state of the art language models contain billions of parameters, and must be trained on massive datasets. As a reference point, we’ll use OpenAI’s GPT-3, since there are some reasonably well grounded numbers available online.

Just the pretraining phase of GPT-3 relied on approximately 570GB of text content, pulled from sources such as Common Crawl, OpenWebText2 and Wikipedia. Prior to training, this content is converted to tokens through a process called tokenization. The resulting training set contains approximately 300 billion tokens.

The largest GPT-3 model uses 175 billion parameters. This translates to roughly 350GB if you store the weights as 16-bit floating points (FP16). This poses many architectural challenges, even when working with data-center class GPUs. It is currently impossible to load so much data onto a single GPU. Therefore, the weights for Large Language Models must be distributed across multi-GPU clusters, with the clusters themselves connected via high-bandwidth interconnects.

Hardware Constraints

So what do we do when we want to run a Large Language Model on an embedded device?

We face a number of challenges when targeting an embedded environment. First is limited memory (RAM/VRAM). LLMs use a lot of RAM for model weights and context windows. Besides limited RAM, the memory architecture may not be as fast as desktop / server RAM.

Next is limited computing power. While many embedded devices now feature NPUs (Neural Processing Units) these are still very limited compared to desktop and server-grade GPUs.

Model Architecture Support

Finally, model architecture support may be limited. Embedded NPUs are simpler in nature than desktop and cloud GPUs. They often lack some of the hardware-accelerated operations required to implement a model using the Transformer architecture. Those operations need to be implemented using slower CPU-based emulation.

Attention kernels and rotary embeddings are two examples that are often not supported.

Transformer Architecture

Understanding the Transformer Architecture will get us a little closer to the question of why language models are so large in the first place. We’ll briefly touch upon the concepts of Self-Attention, Cross-Attention, and the model structures typically seen in LLMs.

Image Credit: Figure 1 from Attention Is All You Need.

Self-Attention

Self-Attention is a mechanism that allows each token in a sequence to decide which other tokens are relevant to it and to aggregate that information via a weighted sum. The result is a context-dependent token representation, which can “look” backwards as well as forwards in the input sequence.

This is in contrast to RNNs and LSTMs—these models could pass context forward from one state to the next; however, they were restricted to encoding that information in hidden state that was of a fixed size. They also suffered from vanishing gradients, and were unidirectional in nature.

Scaled Dot-Product Attention

We hand-waved an important detail above: the weighted sum.

Each token’s initial representation is projected into query, key, and value vectors. The query vector is used to calculate a similarity score with the keys for every other token in the sequence. We then apply a softmax to the similarity scores to find attention weights. Those weights are used to calculate a weighted sum of the corresponding value vectors. The final result is a context-dependent representation of a token. This whole process is called Scaled Dot-Product Attention.

\[\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_k}} \right) V\]

That paragraph was a lot to digest! If this is new to you, I highly recommend the Neural Networks course from 3blue1brown. Specifically, the content on Attention in Transformers.

Cross-Attention

Scaled Dot-Product Attention is also used in Cross-Attention. Where Self-Attention asks “How do tokens in this sequence relate to one other?”, Cross-Attention asks “How do tokens in this sequence relate to tokens in another sequence?”

Both Self-Attention and Cross-Attention use the Scaled Dot-Product Attention formula that we saw above. The only difference is how we construct Q, K, and V.

In Self-Attention, Q, K, and V are all derived from the same sequence of hidden states. You can think of this as trying to predict the next token in a single sequence.

In Cross-Attention, Q comes from one sequence, while K and V come from another, distinct sequence. So you can think of this as trying to predict the next token in an output sequence, given the tokens seen in a different input sequence.

Encoders and Decoders

Cross-Attention was originally developed to support Machine Translation and Text Summarisation tasks. In these architectures, one half of the model is known as the Decoder, and the other half is the Encoder. Cross-Attention allows the decoder to attend to outputs of the encoder, while generating a distinct sequence of tokens. This allows the output to reflect the meaning of the original input, but transformed in some meaningful way - e.g. translated to a different output language.

The basic structure of an Encoder-Decoder Architecture is shown below:

It’s also possible to have encoder-only or decoder-only models. Google’s BERT is an example of an encoder-only model, while GPT is an example of a decoder-only model. In both cases, only Self-Attention is used.

The key difference between GPT (decoder-only) and BERT (encoder-only) lies in how they are trained. BERT is trained using a bidirectional approach, whereby tokens are permitted to attend to any other token in the input text. GPT is unidirectional, and uses masking techniques to ensure that the model cannot “look into the future”.

Model Structure

Transformer models tend to follow a similar architecture.

The first part of the model is the tokenizer—this is the part of the network that splits the input into individual tokens. We explore this in more detail in the Tokenization section below.

Following the tokenizer is the backbone. This is the core of the neural network, implementing the attention mechanism used by the model to generate an output sequence. However, the output of this layer cannot be used as-is. It’s a highly specialised representation of the model’s hidden state.

The last part of the model is the head (sometimes called the “LM head”). This is a standard linear layer that projects each hidden state to a vector of size vocab_size. This vector is used to choose tokens (essentially the reverse of tokenisation), based on the confidence assigned to each possible token.

Recommended Resources

We’ve skipped over many other details, such as Positional Encoding, which are important to understand when diving into LLMs. To learn more, I recommend picking up a copy of Hands-On Large Language Models by Jay Alammar and Maarten Grootendorst.

Hands-on Large Language Models covers a wide range of topics. Beyond just LLM fundamentals, it explores different ways to chain models together to build sophisticated LLM-powered applications.

Tiny Language Models

With Large Language Models out of the way, we can shift our attention to Tiny Language Models.

Tiny Language Models are typically trained using fewer parameters or other simplifications such that they will fit within the constraints of an embedded device. We’ll look at just a few here, but be aware there are many others!

Model Distillation

A key ingredient in creating Tiny Language Models is model distillation. This is the process by which a compact “student” model can be trained using the outputs of a much larger “teacher” model. This allows the student model to capture concepts learned by the teacher, without direct exposure to the training set.

TinyLlama

TinyLlama is a compact language model that copies the architecture and tokenizer of Llama 2 (Meta’s open-source language model) but shrinks the parameter count to make it compatible with modest GPUs, CPUs, and edge devices. With just 1.1B parameters, TinyLLama is approximately 150x smaller than GPT-3.

Because TinyLlama is architecturally compatible with Llama 2, it can be used with many existing Llama-based tools and projects, and there are multiple downstream variants mirroring those of the base model. For example, the TinyLlama-1.1B-Chat variant is fine-tuned on the UltraChat dataset.

Qwen

Qwen (also known as Tongyi Qianwen) is a family of large language models developed by Alibaba Cloud. Qwen is released in a variety of model sizes (all published on Hugging Face), so it’s likely that you’ll find one that fits within your constraints.

At time of writing, the latest iteration is Qwen 3. However, earlier iterations are often preferred for embedded AI, due to their smaller size and simpler architectures. Another advantage of Qwen is that it is released with open weights and a permissive license, and it sits in a sweet spot for “tiny” and “mid-size” local models.

The underlying architecture also works well with quantization schemes. Quantization reduces the precision of model weights and activations, making it easier to fit within the constraints of an embedded device. We’ll look at a Qwen variant in the example at the end of this post.

TinyBERT and DistilBERT

TinyBERT and DistilBERT are both distilled versions of Google’s 2018 language model, BERT.

DistilBERT (released by Hugging Face in 2019) is basically a smaller BERT (66M parameters) that you can fine-tune like normal, with no custom distillation pipeline required. What this means is that you can fine-tune the smaller model directly. In practice, it may perform slightly worse than full BERT, but it can be good enough for many applications while offering better speed and reduced memory usage.

TinyBERT, from Huawei’s Noah’s Ark Lab, uses a two-stage distillation approach (first pass is general, then task-specific). It comes in 4-layer and 6-layer versions. The 4-layer variant is even smaller than DistilBERT (just 14.5M parameters) and considerably faster while offering similar accuracy. The 6-layer variants typically use 67M parameters, depending on configuration.

DistilBERT is often preferable unless you rely on the very small footprint offered by the TinyBERT 4-Layer variation. TinyBERT’s more aggressive distillation process may come at the cost of some accuracy. However, this may be a worthwhile trade-off when building for Edge AI devices.

Tokenization

Tokenization is the very first step in processing text input - it is the process by which raw text input is broken down into smaller units that can be handled by a language model. This critical first step reduces the length of the sequence that a model has to process, while retaining important context.

Tokenization is a deceptively broad topic, and although it’s a critical component in modern SOTA language models, we won’t cover it in great detail here. It really deserves its own post!

Kinds of Tokenization

There is, unfortunately, no one-size-fits-all approach to tokenization. First of all, we have determine whether we want tokens that represent characters, words, or sub-words:

Some models even operate on entire sentences - in this case we would break a text into individual sentences, which are processed in their entirety. This is generally NOT relevant these days.

While sub-word tokenization is currently the most common approach used in SOTA models, character-based (or tokenization-free) models have begun to gain traction. We’ll skip of tokenization-free models here, and focus on Sub-Word Tokenization.

Sub-Word Tokenization

In Sub-Word Tokenization, we break text into smaller units that are larger than characters but smaller than full words.

Key algorithms are Byte-Pair Encoding (BPE), WordPiece and Unigram.

Byte-Pair Encoding

BPE starts with a vocabulary of individual characters and iteratively merges the most frequent pair of symbols until a desired vocabulary size is reached. Merges based on the frequency before adding to the vocabulary.

Let’s look at how BPE would be used to choose a set of tokens for the string AAABDAAABAC:

In the first step, we have just one token ID per character. We then calculate the frequency of each pair of adjacent tokens across the training set. The most frequent pair is merged, and given its own token ID. The training set is updated to use the new token. This process repeats until the desired number of tokens has been generated.

WordPiece

WordPiece is similar to BPE in that it iteratively merges subword units, starting from characters. However, instead of choosing the most frequent pair like BPE, WordPiece uses a probability-based score that approximates how much a new merged symbol would improve the likelihood on the training data.

Unigram

Unigram starts with a large vocabulary of potential sub-words and iteratively removes the least likely ones until the optimal set is found.

You may also see SentencePiece mentioned in the context of tokenization. SentencePiece is an unsupervised text tokenizer and de-tokenizer that implements Byte-Pair Encoding (BPE) and the Unigram language model.

De-Tokenization

The output of a language model often depends on the inverse of tokenization. The LM Head of the model is responsible for converting an internal representation of a sentence or document into a series of tokens, often using the same mapping as the original model input. Those tokens are then converted into human-readable text via de-tokenization.

De-tokenization may be a probabilistic process. If the language model outputs a confidence level for each possible token, then tokens may be sampled according to confidence assigned to them. This can have a huge impact on the quality of the output, and several different approaches are used.

Greedy Decoding

If we want the fastest output, and can tolerate some mistakes, then the simplest option is Greedy Decoding. We always choose the token with the highest confidence. However, these days Beam Search is more popular.

Beam Search

Beam Search aims to strike a balance between Greedy Decoding and an exhaustive search for the best possible output sequence. The idea is that we keep a certain number of partial solutions in memory, determined by the beam size K. At the first time step, we choose the K most probable tokens, based on the logits generated by the model.

On all subsequent steps, we extended each candidate sequence by one token, and re-evaluate the joint probability of that particular sequence occurring. All but the top-K sequences are discarded, and then the process continues. I highly recommend reading Chapter 10.8 - Beam Search from Dive Into Deep Learning to learn more about the inner workings of Beam Search.

RKLLM Toolkit

RKLLM-Toolkit is a software development kit published by Rockchip, that supports model conversion and quantization on a desktop PC. It also includes a device runtime with C/C++ APIs for Rockchip NPUs, making it possible to deploy language models to Rockchip devices.

Image Credit: RKLLM-Toolkit repository on GitHub.

Getting started with RKLLM

Khadas makes it very easy to get started with LLM demos on the Edge2. While the documentation is a bit light, there isn’t actually that much to it.

Assuming you’re running a Linux distribution on your Edge device, you should find a script named khadas_llm.sh in your path. Running this will install dependencies, clone the rknn-llm repo and prompt you download/builds models.

CMake already installed, skipping installation
Repository exists, skipping clone

======================
System Memory: 15.5G
Available Models:
 1) DeepSeek 1.5B        [Missing]
 2) DeepSeek 7B          [Missing]
 3) Qwen2 2B-VL          [Missing]
 4) ChatGLM3 6B          [Missing]
======================
Download missing models? [y/N]

Unfortunately, the first time I tried this, the models failed to download due to 404 Not Found errors. I eventually found the correct files on dl.khadas.com:

From here, I simply modified the khadas_llm.sh script, making the following replacement:

dl.khadas.com/development -> dl.khadas.com/resources/development
                                           ^^^^^^^^^

After doing this, it was smooth sailing!

Available downloads:
 1) DeepSeek 1.5B
 2) DeepSeek 7B
 3) Qwen2 2B-VL
 4) ChatGLM3 6B
Select model to download: 1
Downloading DeepSeek 1.5B...

The model we want is 1) DeepSeek 1.5B. The full name for this model is DeepSeek-R1-Distill-Qwen-1.5B.

NOTE: the order may be different if you run this script on your own device.

DeepSeek-R1-Distill-Qwen

DeepSeek-R1-Distill-Qwen-* is a series of open-weight reasoning-focused language models where DeepSeek’s larger R1 reasoning model has been distilled into a Qwen 2.5–based backbones. This model is available in 1.5B, 7B, 14B, 32B parameter variations. We use the 1.5B parameter model in the example below.

When you first download the DeepSeek-R1-Distill-Qwen-1.5B model, it will be stored in a generic format such as ONNX. Certain conversion steps need to take place before running on the Rockchip NPU:

Converting the model graph to use Rockchip-compatible operations
Layer fusion, which is a kind of optimisation
Quantization, to reduce the precision of model weights to fit within device constraints

This is called building the model. We covered these in my previous post on Edge AI.

The khadas_llm.sh script can do this for you automatically. Simply follow the prompts to build the model:

Building DeepSeek 1.5B model...
-- Configuring done (0.0s)
-- Generating done (0.0s)
-- Build files have been written to: /home/khadas/rknn-llm/examples/DeepSeek-R1-Distill-Qwen-1.5B_Demo/deploy/build/build_linux_aarch64_Release
[100%] Built target llm_demo
[100%] Built target llm_demo
Install the project...
-- Install configuration: "Release"
-- Up-to-date: /home/khadas/rknn-llm/examples/DeepSeek-R1-Distill-Qwen-1.5B_Demo/deploy/install/demo_Linux_aarch64/./llm_demo
-- Up-to-date: /home/khadas/rknn-llm/examples/DeepSeek-R1-Distill-Qwen-1.5B_Demo/deploy/install/demo_Linux_aarch64/lib/librkllmrt.so

Configure runtime parameters (press Enter for defaults)
Max new tokens (default 2048):

Example Prompt

Finally, we can enter some text to see how the model responds. Let’s try asking how many planets are in the solar system…

user: How many planets are there in the solar system?
robot: <think>
Okay, so I need to figure out how many planets there are in our solar system. Hmm, let me start
by recalling what I know about the solar system. I remember that it's a system of eight planets
orbiting the sun. But wait, isn't there something called Pluto as well? I think it was considered
a planet for a long time but later it was classified as a dwarf planet because its size is too
big and it doesn't have an atmosphere.

So, if there are eight planets in total, including Pluto, then excluding Pluto, how many would be
left? Let me count: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune. That's
eight. But wait, I think sometimes people include Pluto as a planet too, so maybe the count is
nine.

I should also consider if there are any other planets besides these eight. From what I know, our
solar system has eight planets: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune.
There aren't any others in our galaxy or beyond it that we can observe as planets right now. So,
excluding dwarf planets like Pluto, there are eight.

I wonder if there's a different way to count them. Maybe including some gas giants or something
else? No, I don't think so. Each planet is distinct and has its own orbit around the sun without
overlapping with others in terms of size or distance from it. So, each one is unique in their
characteristics.

Wait, could there be any other planets that are not considered planets by modern definitions?
I'm pretty sure all eight are recognized as planets now. There's also the asteroid belt and the
Kuiper Belt, but those aren't planets themselves; they're regions of space with many asteroids.

So, putting it all together, if we count each planet individually, there are eight in our solar
system. If we include Pluto, then it's nine. But since Pluto is no longer recognized as a planet,
I think the correct answer is eight.
</think>

There are 8 planets in our solar system. These are Mercury, Venus, Earth, Mars, Jupiter, Saturn,
Uranus, and Neptune. Pluto was once considered a planet but is now classified as a dwarf planet
due to its size and lack of atmosphere.

**Answer:** There are 8 planets in the solar system.

--------------------------------------------------------------------------------------
 Stage         Total Time (ms)  Tokens    Time per Token (ms)      Tokens per Second
--------------------------------------------------------------------------------------
 Prefill       318.77           14        22.77                    43.92
 Generate      33152.13         482       69.91                    14.30
--------------------------------------------------------------------------------------

Next Steps…

Machine Translation

This content in this post started as just one section of a deep-dive into Machine Translation at the Edge, which looks at implementing a machine translation system on the RK3588 NPU. Keep an eye out for that post!

Tokenization Deep Dive

While writing this, I also found myself writing at length about different tokenisation methods, and the latest developments in token-free language models. That content has also been moved to a separate draft, and I hope to have both posts published by the end of the year.

References

As usual, I drew upon dozens of resources while preparing this post. The list below contains a selection of links that I think will be useful to someone wanting to dive deeper into the topics covered in this post. They’re listed roughly in the same order as the topics were covered.

Word Embeddings

The Illustrated Word2Vec

Attention

3blue1brown Neural Networks Course

GPT

Language Models are Unsupervised Multitask Learners - This is the paper that introduced GPT-2
Language Models are Few-Shot Learners - This is the paper that introduced GPT-3

Tokenization

Hugging Face - Unigram Tokenization
Mastering Tokenization in NLP: The Ultimate Guide to Unigram and Beyond!
Let’s build the GPT Tokenizer - Andrei Karpathy’s talk on BPE

Text Generation

Dive Into Deep Learning - Beam Search

Edge AI using the Rockchip NPU

Sun, 20 Jul 2025 00:00:00 +0000

Ever wondered what it takes to run real-time object detection and face recognition on a tiny device, without the cloud? I’ve been exploring this recently, while testing the limits of the Khadas Edge2. Powered by a Rockchip RK3588 processor - with a built in neural processing unit - this little Edge AI device packs quite a punch. And it’s all neatly packaged for consumers and hobbyists alike.

This post is split into three parts. Part 1 will begin with a brief tour of the Edge AI problem space, the kinds of AI models that are available, and technical details of the Edge2. In Part 2, we’ll dive into the software and example code available for Rockchip devices, and explore some popular computer vision datasets. Finally, Part 3 will make this more practical, walking through the steps to train a custom object detector.

Part 1: Edge AI
Khadas Edge2
Low Level Details
Applications
Part 2: Software and Datasets
Khadas Resources
- Oowow
- Edge2 Examples
Quantization
Datasets
Part 3: Transfer Learning
- Feature Extraction
- Fine-Tuning
YOLOv5
Building a Dataset
- FiftyOne
- Data Preparation
Training the Model
- Evaluation
- Conversion
Closing Thoughts

Part 1: Edge AI

Let’s start from the top.

You may already be familiar with Cloud AI, thanks to services such as ChatGPT, or even Google Search, with its introduction of AI generated content into search results. If Cloud AI is one end of the AI spectrum, this post focuses on the other end of that spectrum - machine learning algorithms running on local hardware devices. This is generally refered to as Edge AI, or AI at the Edge. Naturally, there is an appropriately named O’Reilly book on the topic:

To learn more about the book, visit the O’Reilly website.

Some commonplace examples of Edge AI devices include modern mobile phones, smart appliances, and cars with self-driving or driver-assist features. These are cases where a reliable internet connection may not be available, or safety concerns require that AI processing happens with low latency.

Example Applications

Edge AI isn’t limited to the examples above. You’ll also find AI in the following embedded settings:

Smart Cameras: Detecting intruders or anomalies without needing to stream footage to the cloud.
Voice Control: Processing voice commands locally (e.g., wake word detection).
Manufacturing: Identifying defects on the production line.
Retail: Monitoring customer foot traffic and shelf inventory in real time.
Healthcare: Tracking bio-signals locally for faster diagnosis.

With all these examples, it’s not difficult to find some common themes as to why Edge AI can be desireable:

Being able to react to data or signals in real-time, without needing to send data to the cloud. Lower latency means faster, more responsive UX.
Offline operation, independent of an internet connection. This means Edge AI works even when you’re on a flight, or stuck in a tunnel.
Preservation of privacy. Personal data never leaves the device, which can be critical for privacy-focused users, or heavily regulated industries such as healthcare.

Constraints

A key challenge with Edge AI is that devices often have constrained compute, memory, and power budgets that inherently restrict model complexity. This is especially true in the Embedded Systems space. In the most extreme cases, Edge AI may need to run on microprocessor devices such as the Arduino. Of course, it’s common these days to use more powerful alternatives, such as the Arduino Portenta H7:

Image Credit: The Arduino store.

The Portenta H7 features an STM32H747 processor, with two cores - a Cortex-M7 and a Cortex-M4. The separation of these cores means that the device can handle real time control and AI processing simultaneously. However, this processor is not designed for compute-heavy AI tasks such as object detection and real-time video analysis.

Model Updates

The embedded nature of Edge AI devices also leads to difficulties in shipping model updates. Regulatory constraints may require that models are trained and rigorously tested up-front. And once a device is deployed, it may be costly (or even illegal) to update the model. Bandwidth and storage can also be an issue. Some state of the art (SOTA) models may be hundreds of megabytes in size (or even gigabytes), making distribution impractical.

Training vs Inference

Implementation of a machine learning algorithm can be broken down into two key phases: Training and Inference.

During the training phase, a machine learning algorithm is exposed to a dataset from which it can extract (or ‘learn’) patterns and other meaningful relationships. The output of training are the parameters for a model, e.g. connection weights for a neural network. Training tends to be computationally expensive, and requires access to training data. On each pass through the dataset (known as an ‘epoch’), the network predicts the labels associated with each example in the dataset. Those predictions are compared with the true labels, and the error is distributed across the weights in the network:

In the inference phase, the parameters (or weights) learned in the training phase are loaded into memory, and used by the model to process previously unseen inputs. The weights are rarely altered at runtime.

This storage cost for training data is prohibitive for Edge AI applications, and the data itself may be proprietary. As such, Edge AI applications typically focus on efficiency for inference, rather than training.

Khadas Edge2

This takes us to the Khadas Edge2.

The Edge2 is a slim and compact ARM PC that is designed with Pogo Pads, FPC headers and MIPI-Interfaces which makes it convenient for developers to access the necessary I/O for embedded applications. It includes HDMI output, two USB-A ports, and two USB-C ports, for data and power:

The MIPI ports on the bottom edge of the board allow three cameras to be connected directly to the device, for advanced video processing tasks.

Rockchip RK3588

The chip that powers all of this is the Rockchip RK3588S2, an 8-core 64-bit ARM processor. This includes a dedicated NPU (neural processing unit), making it possible to run deep learning inference on-device, without needing a GPU or access to a Cloud AI service. This makes the RK3588 ideal for real world Edge AI applications such as surveillance, robotics, or smart sensors.

The Rockchip NPU has direct access to system memory, and supports INT8, INT16 and FP16 precision models natively. We’ll discuss the importance of precision further in the Numeric Precision section below.

Power Consumption

A critical characteristic of embedded devices is power consumption (along with it’s first cousin, heat dissipation). Embedded devices are often battery-powered, or deployed in locations where it is difficult to provide active cooling.

The Edge2 balances strong performance with very modest power consumption. While under load, Thermal Design Power (TDP) is estimated to be between 5-10W. This is especially appealing when compared to using a desktop GPU, which can easily have a TDP 30-50x greater. The NPU can handle up to 6 TOPS (Trillions of Operations Per Second) - very impressive given the power budget.

Khadas’ documentation recommends a 24W power supply, which will provide ample headroom for powering connected peripherals.

Community Support

On the software side, Rockchip supports several options for deploying models to the device - these include RKNN Toolkit, ONNX, and LiteRT (“Lite Runtime”, formerly known as TensorFlow Lite). RKNN Toolkit provides tooling for converting various model formats into RKNN Format, which run natively on the RK3588 NPU.

Image credit: RKNN Toolkit documentation.

This means you can build and train models in popular frameworks (e.g. PyTorch or TensorFlow) before converting them to run natively on the Rockchip NPU. Rockchip maintains an extensive RKNN Model Zoo on GitHub, which includes example code for many popular models.

Finally, Khadas provides extensive documentation for working with the Edge2 itself, and growing community interest around Rockchip devices on Reddit - see the r/RockchipNPU subreddit.

Compared to NVIDIA Jetson

Many developers are discovering that these RK3588 boards compare very favourably against NVIDIA Jetson or Raspberry Pi for Edge AI inference. The Khadas Edge2 offers competitive edge AI performance at a lower price point than NVIDIA Jetson Xavier NX or Orin.

Multiple I/O options, including M.2, USB-C, GPIO, and HDMI, make the Edge2 suitable for integrating into smart cameras, drones, or custom enclosures.

Low Level Details

Let’s dive deeper into the capabilities of the device, and how inference on an embedded device is different to other environments.

Floating Point

When talking about floating points numbers, or ‘floats’, we usually refer to decimal numbers that are represented using 32-bits. These bits can be broken down into three parts: a sign, an exponent, and a fraction (the ‘mantissa’). This is in accordance with the IEEE 754 standard, which describes how these bits can be interpreted as a decimal.

It’s beyond the scope of this post to explain how all this works. The most intuitive explanation I’ve seen for this is Fabien Sanglard’s Floating Point Visually Explained. And for much more detail, I recommend What Every Computer Scientist Should Know About Floating-Point Arithmetic.

Numeric Precision

The RK3588 SoC has a dedicated NPU that natively supports models with INT8, INT16 and FP16 numeric precision. These refer to 8-bit integer, 16-bit integer, and 16-bit floating point precision, respectively. These are pretty self-explanatory - INT8 is a signed integer, represented using 8 bits. This means it can take values between -128 and 127. Naturally, INT16 is 16 bits, covering a range between -32768 and 32767.

There are several numeric representations commonly used in Machine Learning:

This determines the numerical precision of the weights and activations used during inference, but it is generally NOT the same as the precision used during training. What usually happens is that a model is trained at a higher level of precision (e.g. FP32 or 32-bit floating point) on a desktop or cloud GPU, before being converted for use on the target device.

Memory Architecture

This is another area where the Edge2 differs from a desktop GPU. The Edge2 comes with up to 16GB of LPDDR4 RAM, which is shared by the CPU, GPU, and NPU. This shared memory architecture means data can be easily shared between different processing units, without the overhead of transferring memory from one device to another.

By comparison, the RAM for an NVIDIA desktop GPU is independent of main system memory. While NVIDIA offers a ‘Unified Memory’ solution that allows system memory and GPU memory to share the same address space, this is simply a convenience for the programmer - it does not avoid the need to transfer data between physical devices.

Applications

Now that we know a bit about the kind of device we’re working with, what are some applications that are well suited to it?

Object Detection

Object Detection is a classic computer vision task that involves identifying, classifying and locating objects within an image or video feed. The key questions that object detection algorithms (attempt to) answer are:

How many objects are present? (Identification)
What kinds of object are they? (Classification)
Where are they are located? (Localisation)

Image Credit: Output from Khadas example code for SSD Object Detection

Some common models for Object Detection include SSD (Single Shot Detector), RetinaNet, and Faster R-CNN. Another popular choice is YOLO (short for “You Only Look Once”), which we’ll examine in more detail in the second part of this post.

Facial Recognition

Facial Recognition is a classic computer vision problem, generally broken down into two distinct tasks:

The first is Face Detection, where we would like to locate all of the faces in an image, similar to how Object Detection is primarily concerned with locating objects.

The second is Face Recognition, which has a range of additional applications, such as providing biometric security for personal devices. As you would expect, there are many models available. Just to list a few:

FaceNet (Google, 2015) - This model introduced the concept of face embeddings using triplet loss. The model is trained using example ‘triplets’, each consisting of three separate face images. FaceNet provides a good performance baseline, however newer models may offer superior embedded performance.

Triplet Learning: Positive examples move closer to the anchor, negative examples move further away.
RetinaFace - RetinaFace builds upon the RetinaNet object detection model, adding the ability to detect human faces and predict key facial landmarks (eyes, nose, corners of the mouth).

See the original paper RetinaFace: Single-stage Dense Face Localisation in the Wild for more detail.
ArcFace or MagFace - ArcFace and MagFace are two newer state-of-the-art models for Face Recognition.

ArcFace (2018) introduces a variation of the Softmax Loss function (often referred to referred to as Cross Entropy Loss), which encourages same-class samples to have embeddings closer to one another, while different-class samples are further apart. What this means in practice is tighter clustering of the same person’s face, and clearer separation between different faces.

MagFace (2021) builds on ArcFace by incorporating the magnitude of each embedding to the quality of the input face image. Higher-quality face images yield larger embedding magintudes. This allows MagFace to not only maintain strong recognition performance but also provide a built-in measure of face quality, which can be valuable in real-world applications.

Sentiment Analysis

Sentiment Analysis typically refers to the identification of tone, attitude or emotion present in text or speech. For example, in an e-commerce setting, we could classify online reviews as positive, neutral or negative. For audio or speech, Sentiment Analysis can be used to detect emotional states (e.g. delight, stress or frustration).

When analysing recordings of speech, sentiment can be inferred from acoustic features. These include:

Pitch
Loudness
Tone of voice
Pauses and rhythm of speech

These may be derived from a Mel Spectrogram:

Image Credit: Spectrogram of the spoken words “nineteenth century” from Wikipedia.

A Mel Spectrogram has the frequency axis converted to the mel scale, which is designed to match human auditory perception.

Analysis of text tends to be more subtle, as it may not capture emotional nuances. This can be improved by adding more context. Alternatively, more samples may be required to accurately determine the tone of a body of text.

Part 2: Software and Datasets

Now that we have the lay of the land, let’s see what tools are available for the Edge2 specifically, starting with libraries and tooling from Rockchip.

RKNN Toolkit

RKNN-Toolkit2 is an SDK from Rockchip that supports model conversion, inference and performance evaluation on PC and Rockchip NPU platforms. On the PC side, it supports model conversion and quantization workflows. On the device side, it facilitates loading and running models on the Rockchip NPU, using C/C++ or Python APIs.

The primary purpose of the toolkit is to take models that have been trained on PC, and convert them into a format that will work with the the RKNN API. Once a model has been concerted to RKNN format, we run inference on a Rockchip NPU using the RKNN C API or Python API. The toolkit contains several examples written in both C++ and Python.

It’s worth noting that RKNN models are sensitive to input tensor shapes, memory alignment, and batch sizes. If unsupported features are used, then computation may fall back to the CPU. The easiest way to mitigate this is to use Rockchip’s example code, which is designed with these constraints in mind.

Model Zoo

When it comes to example code, Rockchip also provides a Model Zoo. “A model what?” you may be asking.

Despite the strange name, this is a simple idea. A Model Zoo is a collection of pretrained machine learning models. These models are trained on standard datasets, and are usually provided for popular frameworks (e.g. PyTorch, TensorFlow, or ONNX).

There’s a fun way to visualise what a Model Zoo offers, courtesy of The Asimov Institute. In response to the Cambrian explosion of neural network architectures, they created a poster called the Neural Network Zoo:

Rockchip has published their own RKNN Model Zoo on GitHub, which captures this in a more practical form. This repo contains pretrained models for a variety of common applications. These models are provided in ONNX format, and can be readily converted to device-specific RKNN weights.

Large Language Models

Support for Large Language Models (LLMs) has been broken out into a separate module called RKNN-LLM. While LLMs are less common in Edge AI applications, this is without doubt the fastest growing area of AI, and deploying LLMs in embedded contexts is an area of active development.

LLMs are still very constrained on devices like the RK3588. RKNN-LLM is mostly in early development and supports small-scale models (e.g., TinyLlama, Qwen-1.8B-Int4). Performance is a far cry from GPT-level performance.

Khadas Resources

Khadas OOWOW

The Edge2 comes built-in with Khadas’ OOWOW service. OOWOW, also known as ‘endless wow’ or infinity wow, is a standalone embedded service for cloud OS installation, device maintenance, and more.

OOWOW will be activated automatically if your device storage is empty, otherwise it can be activated by holding the FUNCTION button while pressing RESET. Once you have configured network access, you will be able to install one of several Android or Linux-based distribution from the cloud!

OOWOW Image Selection

Installation only takes a few minutes. Khadas provides documentation and sample code for each platform (Android and Linux), so this is a quick and easy way to get started with the Edge2.

Edge2 Examples

Khadas provides a range of C++ and Python examples for the Edge2, which can be found in the edge2-npu repo on GitHub. This includes C++ and Python examples. The examples all depend on librknn, Rockchip’s custom neural network library. Many of the examples also depend on OpenCV (both C++ and Python).

The C++ examples are much more extensive than those for Python. While there are over ten examples for C++, there are just four for Python.

Quantization

An important part of running a model of the Edge2 is ‘Quantization’ - this is the process by which the precision of model weights is reduced. This can significantly reduce the memory footprint of the model, and ensures that we can take advantage of hardware acceleration available at lower levels of precision. This is typically achieved by applying a ‘scale factor’ that reduces the range of possible values (e.g. -128 to 127 for INT8), then truncating those values to the desired precision.

For example, a model trained with 32-bit floating point (FP32) weights can be quantised to 16-bit floating point (FP16), or even 8-bit integers (INT8), significantly reducing model size and improving inference efficiency, particularly on edge devices.

To learn more about the theory of Quantization, I suggest reading Hugging Face’s Conceptual Guide to Quantization.

Calibration

RKNN Toolkit provides tooling to quantise models while converting them to RKNN format. This not only adjusts the weights, but also adjusts the calculations performed throughout the model to compensate for the scaled down ranges.

Unfortunately, we don’t get this for free. Calibration relies on a selection of sample inputs, which are used to choose the scale factor and zero point values for model weights. The simplest option is Symmetric Quantization, in which the zero point from the original FP range is mapped onto zero in the target (e.g. INT8) range:

Symmetric Quantization: Notice that some distinct FP values map onto the same INT8 value.
There is also a region of wasted INT8 range on the left.

The reason this is suboptimal is that the minimum and maximum values of the FP range must have the same absolute value. To see why this is wasteful, consider a weight that takes on values from -512 to 1024. Although the true minimum is -512, the range must extend all the way to -1024.

Asymmetric Quantization

Asymmetric Quantization addresses this by allowing the zero point to be shifted to the left, or to the right:

Asymmetric Quantization: Zero point is shifted by -85 units. More of the INT8 range is mapped to positive FP values.
No unused INT8 values on the left.

The key difference here is that the zero point can be chosen as the midpoint between the minimum and maximum FP values. The resulting INT8 range will then cover the entire FP range.

Gotchas

Quantization may seem like a great idea, with very few drawbacks. In practice, it can be hard to track down issues with loss of accuracy, or overall performance of the model. RKNN models are sensitive to input tensor shapes, memory alignment, and batch sizes - these can all impact model performance at runtime.

With YOLOv5 for example, the example code includes some small changes to the underlying model graph. Specifically, a post-processing step is moved outside the main model. This is to address stability issues introduced by quantization.

Datasets

Before we talk about the resources that available for/from Rockchip specifically, it’s worth mentioning some of the datasets that are relevant to Edge AI applications.

MNIST and ImageNet

MNIST and ImageNet are two of the most well known and battle-tested datasets in the field of computer vision models. MNIST is a dataset containing 70,000 grayscale examples of handwritten digits:

Image Credit: MNIST Examples courtesy of Wikipedia

Performance on MNIST was an important benchmark for early Convolutional Neural Networks (CNNs). Beginning with the ImageNet Large Scale Visual Recognition Challenge in 2012, ImageNet became the defacto standard for measuring performance on image classification tasks. The dataset used for that competition (ImageNet-1K) contains ~1.2 million images, across 1000 classes.

The full ImageNet dataset contains ~14 million labeled images across ~20,000 categories. Although it was created almost 20 years ago, ImageNet continues to be used for training and benchmarking deep learning models on object detection and image classification tasks.

VGGFace2

VGGFace2 is a large-scale face recognition dataset. It consists of ~3.3 million images of ~9100 subjects, downloaded from Google Image Search. The dataset has been designed to capture the variations that arise due to age, socioeconomic background, and the setting in which the images are captured:

Image credit: Example taken from VGGFace2: A dataset for recognising faces across pose and age.

COCO Dataset

COCO is a large-scale object detection, segmentation, and captioning dataset. It is often used to train and benchmark object detection models, such as YOLO. However the wide range of annotations included in the dataset make it suitable for image segmentation, human pose estimation, and even captioning.

Image credit: Adapted from Figure 1 in Microsoft COCO: Common Objects in Context.

You will often find pretrained object detection models that have been trained on the COCO dataset. COCO defines 91 classes in total, but the dataset only uses 80 of those classes. An off-the-shelf pretrained model will readily categorise objects into one of those 80 classes, and we can use Transfer Learning to re-fit this to our own custom datasets. This is cover in more detail below.

Open Images Dataset

The Open Images Dataset (OID) is another large scale dataset that can be extremely valuable when customising pretrained models. OID contains approximately 9 million images, annotated with image-level labels, object bounding boxes, object segmentation masks, visual relationships, and localised narratives.

The maintainers of OID have partened with Google to offer a tool called FiftyOne that can be used to download subsets of the dataset, as needed for specific tasks. For example, if you wanted to train an object detector to recognise firearms, you could download images (and annotations) for classes such as Handgun, Rifle and Shotgun.

Part 3: Transfer Learning

We now have the fundamental knowledge required to train our own custom object detector. We will take advantage of Transfer Learning, to avoid the cost and complexity of training a model from scratch.

Transfer Learning refers to the practice of using a pretrained model, typical trained on a benchmark dataset like ImageNet or COCO, for a different but related task.

The two main approaches to Transfer Learning are Feature Extraction and Fine-Tuning. These differ in how they treat the existing layers of the model - specifically, whether they are ‘frozen’ or ‘unfrozen’.

Feature Extraction

In Feature Extraction, we freeze all but the top few pretrained layers. Those top layers are then re-trained from scratch, for the new task. This kind of Transfer Learning is very fast, and is well suited to cases where the target task is similar to the original task, or when the target domain is a subset of the original domain.

It’s worth nothing that when training a classifier, we will typically replace the topmost layer entirely. If the top layer previously had 80 outputs (e.g. corresponding to the 80 classes in COCO), we would replace that with a new layer, ensuring that the new layer has the number of outputs required for our custom task.

Fine-Tuning

Fine-Tuning takes this a step further. During training, we unfreeze all of the pretrained layers and train them along with the new layers on a custom dataset. Fine-Tuning is more effective when the target dataset is large enough or differs significantly from the original domain.

In practice, these two approaches are often combined: training begins with pure Feature Extraction, using frozen pretrained weights. Once the model achieves a reasonable level of performance, the pretrained weights are unfrozen, allowing them to adapt to the new task.

YOLOv5

YOLO, or “You Only Look Once”, is best thought of as a family of object detection models. First introduced in 2015 (original paper here), YOLO has continued to be improved and iterated upon by multiple researchers. You can identify these iterations by a version number suffix (e.g. YOLOv5, YOLOv8, and so on). Many recent advancements have been driven by Ultralytics - an AI company that has built its reputation around pushing the envelope with YOLO object detection.

A deep dive into the implementation details of YOLOv5 is beyond the scope of this post. However, if you would like to learn more, a great starting point is Ultralytics YOLOv5 Architecture.

Image credit: A small portion of the YOLOv5 Model Structure from Ultralytics YOLOv5 Architecture.

Source Code

The fastest way to get started with YOLOv5 is via the yolov5 GitHub repo:

# Clone the YOLOv5 repository
git clone https://github.com/ultralytics/yolov5

# Navigate to the cloned directory
cd yolov5

# Install required packages
pip install -r requirements.txt

From here, you can easily run YOLOv5 object detection on your PC. Here are just some of the examples from the yolov5 README file:

# Run inference on a local image file
python detect.py --weights yolov5s.pt --source img.jpg

# Run inference on a local video file
python detect.py --weights yolov5s.pt --source vid.mp4

# Run inference on a screen capture
python detect.py --weights yolov5s.pt --source screen
...

The --weights argument tells detect.py which version of the model to use. There are several to choose from, described below.

There are other scripts in the yolov5 repo that will be of interest to us:

export.py - Exports a YOLOv5 PyTorch model to other formats, e.g. ONNX.
train.py - Trains a YOLOv5 model on a custom dataset.

Note that export.py can not export directly to RKNN format. For that, we’ll first need to export a model to ONNX format, then convert that to RKNN format using RKNN Toolkit.

Pretrained Models

Ultralytics provides several pretrained models for YOLOv5. Each is a slightly different variation on the model - larger models have wider layers. The number of parameters is proportional to the Width-Depth Multiple of the model, which scales the number of layers in the model, as well as the width of those layers:

 Model              Width-Depth Multiple   Parameters   Size (MB)   Speed      mAP@0.5 (COCO val)
------------------+----------------------+------------+-----------+----------+-------------------
 yolov5n (nano)     0.25x                       ~1.9M     ~3.9 MB   fastest    ~28.0
 yolov5s (small)    0.33x                       ~7.5M    ~14.0 MB   fast       ~37.4
 yolov5m (medium)   0.67x                      ~21.8M    ~41.0 MB   medium     ~45.4
 yolov5l (large)    1.00x                      ~46.5M    ~92.0 MB   slower     ~49.0
 yolov5x (xlarge)   1.25x                      ~87.7M   ~166.0 MB   slowest    ~50.7

The nano and small variations are well suited to rapid iteration in Edge AI and embedded contexts, where memory and compute are at a premium. However as you can see from the table, this comes at the cost of mAP, or Mean Average Precision. It’s worth digging into this in some detail.

Performance Metrics

For object detection models, the metrics we most care about are Mean Average Precision (mAP) and Intersection over Union (IoU). IoU is best understood in graphical form. Given a ground truth bounding box, and a bounding box prediction, we first calculate the area of the intersection and union of those boxes. We then take the ratio of those two values:

Mean Average Precision is an aggregated metric, that takes IoU into account. Precision answers the question: “Given all the objects the model identified, what proportion of these were actually correct?”

In the results shown above, mAP50 is the Mean Average Precision at IoU = 0.50. This means a prediction is considered correct if its IoU with a ground truth box is at least 0.50.

This measurement can be made stricter and more comprehensive, by averaging the mAP over a range of IoU thresholds. The mAP50-95 metric is the Mean Average Precision averaged over IoU thresholds from 0.50 to 0.95. This is calculated in steps of 0.05 (i.e. 0.50, 0.55, …, 0.90, 0.95).

Head, Neck and Backbone

An important detail of the YOLOv5 model is that it can be separated into three parts:

The Backbone - This is the main body of the model, responsible for extracting details such as patterns and geometric features.
The Neck - This takes the information gathered by the backbone, reducing it down by combining and aggregating features.
The Head - This uses the information aggregated by the neck to make predictions about objects in the image.

This separation allows us to use Transfer Learning to adapt the model to new tasks. We can take advantage of the pretrained Backbone and Neck weights by leaving those frozen - this means they won’t change while we train our custom model. The Head will be replaced with new dense layers, specifically for our custom task. It’s those new layers that will be trained on our custom dataset.

In order to do that, we need to construct an appropriate training dataset.

Building a Dataset

To make the most of Transfer Learning, you need a custom dataset with examples that are representative of a new task. These examples must be annotated using YOLO-style annotations for bounding boxes. The YOLO annotation format is very simple:

<class_id> <x_center> <y_center> <width> <height>

All coordinates are normalized (range 0.0 to 1.0), and the class ID is a numeric index. The advantage of normalized coordinates is that they are independent of a particular image size.

FiftyOne

The Open Images Dataset (discussed earlier) is a great resource for building custom datasets. The OID website provides instructions on how to download a subset of images and corresponding annotations for your task. This includes an open source tool called FiftyOne from Voxel51:

Image Credit: Screenshot from the introductory video included in the FiftyOne Docs

One you have created a custom dataset, you should split it into training, validation and test sets. The training dataset will be used to train the model, while the validation dataset will be used to tweak hyperparameters and to choose a base pretrained model. The test set can be used at the end to evaluate overall performance.

Data Preparation

Since we’re training a YOLOv5 model, we can use train.py from the yolov5 GitHub repo. This distribution contains the latest and greatest that Ultralytics has to offer. To train on a custom dataset, we use the --data argument to provide the path to a YAML file that describes our dataset:

python train.py --img 640 --batch 16 --epochs 20 --data datasets/firearms.yaml

The other arguments here specify the input image size, batch size, and number of training epocs to complete.

There isn’t much to the dataset YAML file, but there is one detail that can be confusing at first:

# this path is relative to the train.py script
path: ../datasets/firearms
train: images/train
test: images/test
val: images/validation

# Classes
names:
  0: Handgun
  1: Rifle
  2: Shotgun

The gotcha here is that path is relative to the train.py script. Let’s say you’ve checked out the YOLOv5 source into a directory called yolov5 and the firearms dataset in a datasets directory alongside that. Printing out the directory tree using tree should look something like this:

% tree . -L 3
.
├── datasets
│   ├── firearms
│   │   ├── images
│   │   └── labels
│   └── firearms.yaml
└── yolov5
    ├── benchmarks.py
    ├── CITATION.cff
    ├── classify
    |
    :   // trimmed for brevity //
    |
    ├── train.py
    └── val.py

22 directories, 61 files

When you set your path, you need to take this directory structure into account.

Training the Model

Training with YOLOv5 will print out heaps of interesting information process about the model, and how it is performing after each epoch.

Assuming you run the same train.py command as above:

python yolov5/train.py \
  --img 640 \
  --batch 16 \
  --epochs 50 \
  --data datasets/firearms.yaml

You should get some output that is similar to this (at the end of the output):

Validating yolov5/runs/train/exp12/weights/best.pt...
Fusing layers...
Model summary: 157 layers, 7018216 parameters, 0 gradients, 15.8 GFLOPs
                 Class     Images  Instances          P          R      mAP50   mAP50-95
                   all         99        129      0.656      0.714      0.722      0.502
               Handgun         99         14      0.677      0.929      0.856      0.687
                 Rifle         99         89      0.824      0.674      0.796      0.532
               Shotgun         99         26      0.467      0.538      0.515      0.287
Results saved to yolov5/runs/train/exp12

The output includes mAP@50 scores for each of the classes in the custom dataset (closer to 1 is better), as well as the mAP@50-95 scores. It’s normal for the mAP@50-95 scores to be lower, since they are stricter and more conservative.

The weights of each run are stored in yolov5/run/trains/exp<run>. You can see above that the output is from run 12, since the path is yolov5/runs/train/exp12. After each run, the best weights will be found in weights/best.pt. This model file is in PyTorch format.

Evaluation

After training, you can evaluate your model using the test set to measure its performance:

python yolov5/val.py \
  --task test \
  --weights yolov5/runs/train/exp12/weights/best.pt \
  --data datasets/firearms.yaml \
  --img 640

The results should look like something like this:

Fusing layers...
Model summary: 157 layers, 7018216 parameters, 0 gradients, 15.8 GFLOPs
                 Class     Images  Instances          P          R      mAP50   mAP50-95
                   all        307        382      0.611       0.57      0.581      0.373
               Handgun        307         57      0.665      0.737      0.698      0.483
                 Rifle        307        259      0.753      0.517      0.677      0.436
               Shotgun        307         66      0.415      0.455      0.368        0.2
Speed: 0.2ms pre-process, 2.8ms inference, 1.7ms NMS per image at shape (32, 3, 640, 640)
Results saved to yolov5/runs/val/exp

Notice there is a reduction in mAP50 and mAP50-95 scores. This tells us that the model does not fully generalise to examples outside of the training dataset. This can indicate the model is over-fitting the dataset.

Conversion

The final step in this process is to prepare the model to run on the Edge2. This requires two conversion steps:

First, we must convert the model from PyTorch to ONNX format. ONNX is a neural network interchange format that is supported as an input format by RKNN Toolkit. This is trivial using the export.py script from Ultralytics YOLOv5:

python yolov5/export.py \
  --data datasets/firearms.yaml \
  --img 640 \
  --include onnx \
  --weights yolov5/runs/train/exp12/weights/best.pt

This will write the ONNX model data to yolov5/runs/train/exp12/weights/best.onnx.

Second, we must convert the model to RKNN format, so that it can run natively on a Rockchip NPU. There are many reasons why a model may not run natively on the NPU (e.g., unsupported layer types, tensor shapes, or memory alignment). This is one of the reasons why it is worth exploring the models in the RKNN Model Zoo.

The conversion to RKNN format is a little more involved, since we need to provide a calibration dataset for quantization. The calibration dataset is simply a subset of images from the full training set, which can be used to estimate zero points and scale factors during quantization.

python3 rknn_model_zoo/examples/yolov5/python/convert.py \
  yolov5/runs/train/exp12/weights/best.onnx \
  firearms-calib.txt \
  rk3588

Closing Thoughts

With native support for INT8/FP16 precision, an efficient NPU, and plenty of examples via RKNN Toolkit, the Rockchip RK3588 offers a well-rounded platform for running real-world AI workloads at the edge. This post has shown that it’s very achieveable to fine-tune and deploy YOLOv5 without needing a server, GPU, or internet connection.

Edge AI is not without its challenges and constraints. As you get further into converting and deploying your own models on the Edge2, you will discover that Machine Learning demands a high level of care. It’s important to invest time in understanding the nuances of quantization, transfer learning, dataset preparation, and evaluation. These are all essential components in an end-to-end inference deployment.

Where to from here?

My hope is that this post inspires you to experiment with some of your own ML models - even without an Edge2, most of what we’ve explored in this post applies to desktop ML too! As for my own experiments, I’m beginning to explore Machine Translation at the Edge, and what is achievable with Tiny Language Models. Stay tuned for that!

References

These are just a handful of references I used while writing up this post. There are many others, and I will update this list as I incorporate feedback.

YOLO:

Mini-YOLOv3: Real-Time Object Detector for Embedded Applications

COCO:

Quantization:

Floating Point:

Fabien Sanglard: Floating Point Visually Explained

ImageNet / ILSVRC:

ImageNet Large Scale Visual Recognition Challenge 2012 (ILSVRC2012)

Chromecast Device Authentication

Sat, 22 Mar 2025 00:00:00 +0000

Just recently, there was a high profile Chromecast outage, affecting second generation (Chromecast 2) devices. This happened because The Chromecast 2’s device authentication certificate has expired. This gave me an excuse to revisit some older content that I never got around to publishing, about how one can spoof the Device Authentication process - the same process that failed here.

My initial project resulted in a partially-working Chromecast Receiver, written in Go (code here). This implements enough of the Google Cast protocol to receive screen mirrored content from Chrome. However, this post focuses on a particular aspect of that project, which is how a third party Google Cast implementation can masquerade as a genuine Chromecast receiver.

Introduction
Google Cast Protocol
Device Authentication
Aside: Chromecast Outage
Problem Statement
Root Access
Signature Generation
Conclusion
References

Introduction

If you’ve used a Chromecast before, you may be familiar with the Cast extension in Google Chrome (now built in to the browser). You may have even used the Cast functionality available on an Android device, to mirror your screen or play a video on a Chromecast.

Or perhaps, you’ve never seen a Chromecast. This is what a second generation Chromecast looks like:

While the Chromecast design has been through a number of iterations, the basic concept remains the same. It’s a small USB-powered device, that you connect to your TV via HDMI. Basic capabilities include Wi-Fi, Bluetooth, and hardware accelerated video decoding, all in a very compact package. Chromecast also offers some support third-party apps, which run on a minimal web browser.

Google Cast Protocol

The protocol underlying all of this is called Google Cast. This is Google’s proprietary protocol for launching and controlling applications on a Chromecast device, and for sharing video and audio content to Chromecast-compatible devices. While our main focus in this post is device authentication, we’ll touch upon how the Google Cast protocol has been designed to accomodate custom Chromecast apps.

But first, we should talk about the difference between Chromecast Senders and Chromecast Receivers.

Senders and Receivers

As the name suggests, a Chromecast Receiver serves primarily to receive and play video/audio content. It can also be used to mirror the screen of a Chromecast Sender.

This isn’t all a receiver can do. An official Chromecast also bundles a stripped down web browser, which can be used to display web-based content. This makes the device capable of hosting third-party Applications. This allows the Chromecast to run third-party apps, which are effectively single page apps that can play video/audio from external sources such as YouTube. Third-party apps are typically controlled from a user’s mobile device, with a control channel established for sending messages back and forth.

The Google Cast protocol provides namespace and messaging semantics that make all of this work.

Protobuf

The messages used by the Google Cast protocol are defined using Google’s Protocol Buffers data serialisation format, more commonly known as Protobuf. Protobuf takes a single file that describes the serialisation format for a protocol, and uses that to generate language-specific bindings that marshall and un-marshall data in that format.

Image credit: Protobuf Overview from Protocol Buffers Documentation.

At a high level, each Google Cast message includes the following fields:

Namespace
Source ID
Destination ID
Payload (can be binary or a UTF-8 string)

The payload can be binary or text (assumed to be JSON formatted). The namespace tells us how to interpret the contents of the payload. The source ID and destination ID are included because a single communication channel may carry messages that are intended for different Chromecast applications.

You can see how these are defined in Protobuf format in the file cast_channel.proto from the Chromium source code. Note that this file is now part of the Open Screen Library, which aims to implement the Open Screen Protocol, Multicast DNS, and the Google Cast protocol.

Device Discovery

The next topic we should cover is Device Discovery. This is the mechanism by which a Chromecast Sender (e.g. your web browser or Android device) finds Chromecast Receivers on your local network.

Device Discovery relies on a protocol called Multicast DNS, or mDNS for short. As you may guess from the name, Multicast DNS is a variation on the traditional DNS protocol that powers the internet.

With traditional DNS it is possible for a host to query other services on a network using human-readable domain names (e.g. google.com). For example, when you enter tristanpenman.com into your web browser, DNS is used to find the IP address of the web server that contains the relevant content.

A limitation of DNS is that it is centralised, in the sense that we must know the IP address of a DNS server ahead of time. The DNS server must also know about devices on the network ahead of time, via manual configuration. Unfortunately, this doesn’t work well in local network scenarios where DNS may not be so easy to manage. The Multicast DNS standard is designed to overcome this limitation.

Multicast DNS

Multicast DNS, or mDNS for short, handles service lookup in a de-centralised fashion. Instead of directing queries to a specific DNS server, queries are broadcast on the local network using a multicast address. These queries request information about particular services available on the network. Broadcast packets are forwarded to all devices on the network, and any device is free to respond to any query.

To discover a Chromecast, a client will broadcast a query for devices that offer the _googlecast._tcp service. Chromecast devices on the local network will then broadcast a response listing their capabilities, and how to connect (IP and port). This is called ‘advertising’ a service. mDNS responses are broadcast to all devices on the local network, so it is possible for clients to discover new Chromecast devices passively, even they are not actively querying for them.

At startup, a Chromecast Receiver will advertise the existence of the _googlecast._tcp service, on port 8009. It will also respond to queries for the _googlecast._tcp so that new Chromecast Senders are able to find it when they are first started.

Once a Chromecast Sender has discovered a Receiver, it is free to connect to it on port 8009. Connections are secured using TLS, which is an important detail, as it relates to the Device Authentication mechanism that is part of the Google Cast protocol. We’ll come back to this shortly.

Device Authentication

A key challenge in implementing a viable Chromecast Receiver is Device Authentication. This is the mechanism by which a Chromecast Sender determines that a device is a genuine Chromecast, or an officially licensed Chromecast-compatible Receiver.

The way this works is that a Chromecast Sender issues a Device Authentication challenge immediately after connecting to a Chromecast Receiver. The Receiver is expected to extract the challenge payload, sign it cryptographically using a private key that is unique to the Receiver device, thus producing a valid Device Authentication response.

Google Cast connections are secured using TLS. At a minimum, this requires that a Chromecast receiver has a self-signed x509 certificate. While this certificate is self-signed, and therefore not part of a traditional PKI certificate chain, the Google Cast protocol provides another means by which this certificate can be verified as belonging genuine Chromecast device. That process is the focus of this post.

Motivation

Let’s start with the basics. We may ask: Why bother verifying the authenticity of a Chromecast Peer Certificate?

There are several reasons to do this:

Google would look to maintain control over the Google Cast ecosystem. By verifying that a Peer Certificate belongs to a genuine Chromecast, or a licensed third-party device, Google can ensure that the availability of devices in the Google Cast ecosystem is well regulated.
Users would like to know that the device they are connected to is in fact the device they expect. Without Device Authentication it is possible for an attacker to masquerade as another device on the network, possibly intercepting connections.

In a nutshell, a Device Authentication challenge allows a Chromecast Sender to verify that it is connected to a genuine, or officially licensed, Chromecast device, and that the connection has not been intercepted.

Peer Certificates

So how does this Device Authentication mechanism work?

I mentioned above that Google Cast TLS connections are secured using a self-signed x509 certificate. This certificate, hereafter called the Peer Certificate, is regenerated every 24 hours. On its own, this certificate offers little assurance that a user has connected to a genuine Chromecast device.

We need three things if we are to work around this limitation:

The Chromecast receiver must be able to sign its Peer Certificate, using a private key that is fused in the CPU / SoC. We call this key the Device Private Key, and it has a corresponding Device Certificate, which can be found in the Chromecast file system.
The Chromecast sender must be able to verify that the Peer Certificate was signed using the Device Private Key. It can do this using the public key contained in the Device Certificate.
Finally, the Chromecast Sender must be able to verify the authenticity of the Device Certificate. i.e. that the Device Certificate was generated by an entity that is trusted by Google.

In the case of an official Chromecast, the Device Private Key is fused into the Soc (System on a Chip), and cannot be accessed directly by code running on the device. Instead, the CPU provides cryptographic instructions that allow it to sign a small amount of data (e.g a SHA256 hash) using the private key.

PKI

The second part of the process - verifying the authenticity of the Device Certificate - is achieved using PKI, or Public Key Infrastructure. PKI allows a series of certificates to verified, all the way up to a ‘root’ certificate. Together, the certificates forms a ‘chain of trust’.

In the case of Chromecast, there are three certificates that are included in this chain. The first is the Device Certificate - unique to the Chromecast device. The second certificate corresponds to an Intermediate Certificate Authority (or ICA). Each Chromecast generation has its own ICA. The third certificate is the Google Root CA Certificate.

The complete chain looks like this:

Device Certificate
        |
        |---------------> ICA Certificate
            verified by         |
                                |---------------> Root CA Certificate
                                    verified by

This is a lot to digest, so hopefully a diagram will help…

In this diagram, blue boxes represent information that must be known to the receiver software at runtime. This becomes more relevant when we implement Device Authentication ourselves. We can also see here that the Device Private Key is kept in a secure enclave, which can only be accessed indirectly by cryptographic CPU instructions.

ICA Certificates

The ICA Certificates are of particular interest, because they are shared across an entire product line. At the time that a Chromecast device is manufactured, it’s fused Private Key is signed using the ICA Private Key corresponding to this certificate.

We can use OpenSSL command line tools to inspect the ICA certificate:

Certificate:
    Data:
        Version: 3 (0x2)
        Serial Number: 1 (0x1)
        Signature Algorithm: sha1WithRSAEncryption
        Issuer: C=US, ST=California, L=Mountain View, O=Google Inc, OU=Google TV, CN=Eureka Root CA
        Validity
            Not Before: Dec 19 00:47:12 2012 GMT
            Not After : Dec 14 00:47:12 2032 GMT
        Subject: C=US, ST=California, L=Mountain View, O=Google Inc, OU=Google TV, CN=Eureka Gen1 ICA
        Subject Public Key Info:
            Public Key Algorithm: rsaEncryption
                Public-Key: (2048 bit)
                Modulus:
                    00:bc:22:80:bd:80:f6:3a:21:00:3b:ae:76:5e:35:
                    7f:3d:c3:64:5c:55:94:86:34:2f:05:87:28:cd:f7:
                    69:8c:17:b3:50:a7:b8:82:fa:df:c7:43:2d:d6:7e:
                    ab:a0:6f:b7:13:72:80:a4:47:15:c1:20:99:50:cd:
                    ec:14:62:09:5b:a4:98:cd:d2:41:b6:36:4e:ff:e8:
                    2e:32:30:4a:81:a8:42:a3:6c:9b:33:6e:ca:b2:f5:
                    53:66:e0:27:53:86:1a:85:1e:a7:39:3f:4a:77:8e:
                    fb:54:66:66:fb:58:54:c0:5e:39:c7:f5:50:06:0b:
                    e0:8a:d4:ce:e1:6a:55:1f:8b:17:00:e6:69:a3:27:
                    e6:08:25:69:3c:12:9d:8d:05:2c:d6:2e:a2:31:de:
                    b4:52:50:d6:20:49:de:71:a0:f9:ad:20:40:12:f1:
                    dd:25:eb:d5:e6:b8:36:f4:d6:8f:7f:ca:43:dc:d7:
                    10:5b:e6:3f:51:8a:85:b3:f3:ff:f6:03:2d:cb:23:
                    4f:9c:ad:18:e7:93:05:8c:ac:52:9a:f7:4c:e9:99:
                    7a:be:6e:7e:4d:0a:e3:c6:1c:a9:93:fa:3a:a5:91:
                    5d:1c:bd:66:eb:cc:60:dc:86:74:ca:cf:f8:92:1c:
                    98:7d:57:fa:61:47:9e:ab:80:b7:e4:48:80:2a:92:
                    c5:1b
                Exponent: 65537 (0x10001)
        X509v3 extensions:
            X509v3 Basic Constraints:
                CA:TRUE, pathlen:1
    Signature Algorithm: sha1WithRSAEncryption
    Signature Value:
        8b:d4:a1:b1:cf:5d:cd:7b:6c:48:4a:41:1f:53:2f:18:2d:32:
        45:ff:9e:ab:d3:73:3e:1f:22:d7:ea:fa:01:e6:73:03:0f:2b:
        c6:25:bb:a5:ee:c5:f5:45:cb:24:12:2a:ad:c2:5d:05:f4:7a:
        f5:c2:9b:10:16:5a:d1:0a:73:c5:16:39:a0:10:ca:d1:68:85:
        9e:fb:9e:26:83:8e:58:f3:77:a0:4e:e5:db:97:be:2d:00:5f:
        f5:94:db:b1:9d:65:6b:fd:f0:d1:04:51:df:cc:92:a6:99:2d:
        71:f5:4d:d5:23:fe:33:1c:a9:b4:ab:c5:bf:1a:b8:d1:80:ef:
        89:c9:e2:1f:9c:4c:48:3b:a2:fa:02:0a:dc:84:01:8a:87:02:
        fb:59:ee:a7:4c:04:7d:74:99:87:6a:25:44:ad:16:aa:ec:4e:
        35:1b:7c:7b:84:c9:b1:3f:e1:82:70:e5:0d:e7:d9:6d:fa:95:
        b6:c5:e4:1e:e8:11:9b:d8:b2:f3:a4:fd:13:f3:83:4f:f7:07:
        14:20:bb:22:a5:a6:8f:d6:b5:db:a9:74:78:e2:93:0d:e5:23:
        2f:05:17:e0:b2:97:67:34:4d:0f:9c:76:43:7b:a6:21:4a:56:
        05:f6:2a:7c:f2:7f:12:94:82:26:29:07:f0:0b:6c:6c:79:14:
        b0:74:d5:6c

This is the same certificate as viewed in Apple’s Keychain app.

Here we can see the Common Name associated with the certificate (Eureka Gen1 ICA) - Eureka being the original codename for Chromecast.

Challenge Request

This gives us the pieces we need to under the Challenge Request.

When a client first connects to a receiver, it may issue a Challenge Request. This is a message that instructs the receiver to sign a hash of its Peer Certificate using the Device Private Key, and to respond with an x509 certificate chain that can be used to verify that this signature is valid.

Here is the overall process:

The client re-hashes the Peer Certificate, then uses the Device Certificate to verify that the signature was generated by that particular receiver. i.e. the signature is verified using the device’s public key contained in the Device Certificate.
The client then verifies that the Device Certificate was signed by an ICA corresponding to the particular model of Chromecast device. In the challenge response, the receiver is expected to provide an ICA Certificate for this purpose.
Finally, the client verifies that the ICA Certificate was signed by a trusted Root CA. For this, it will use a Root CA certificate that is included with Google Chrome or Android OS.

Aside: Chromecast Outage

As an aside, we now have everything we need to know to understand the recent Chromecast outage.

If we simply take a look at the ICA Certificate used for second-gen Chromecast devices - specifically the Validity period - we’ll see that the certificate expired on 9th March 2025:

Certificate:
    Data:
        Version: 3 (0x2)
        Serial Number: 36 (0x24)
        Signature Algorithm: sha256WithRSAEncryption
        Issuer: C=US, ST=California, L=Mountain View, O=Google Inc, OU=Cast, CN=Cast Root CA
        Validity
            Not Before: Mar 12 16:44:39 2015 GMT
            Not After : Mar  9 16:44:39 2025 GMT
        Subject: C=US, ST=California, L=Mountain View, O=Google Inc, OU=Cast, CN=Chromecast ICA 3
        Subject Public Key Info:
            Public Key Algorithm: rsaEncryption
                Public-Key: (2048 bit)
                Modulus:
                    00:d1:de:fb:ad:8b:43:07:28:ae:56:2d:f2:73:2a:
                    1f:63:43:76:6d:8d:b8:d1:d4:90:29:1b:91:68:4a:
                    55:41:a0:d5:61:b4:ec:dd:ae:e1:fa:a7:b6:38:c4:
                    de:19:e1:33:4d:9a:29:f1:48:e2:6b:a7:2c:21:14:
                    22:3f:87:81:f3:71:2c:e6:43:1c:b8:d4:ec:cf:67:
                    2f:b2:a2:75:8b:10:bd:f9:e7:c9:5c:de:05:a9:b4:
                    86:b7:68:7d:a7:76:85:e2:65:b8:76:51:4f:b9:60:
                    5d:7e:2b:64:48:12:66:d9:a7:bb:7c:d7:48:88:8a:

                    <// SNIP //>

The reason it was so difficult for Google to respond to this issue is that second-gen Chromecasts were hard-coded to use this particular certificate.

Problem Statement

At this stage, we can begin to flesh out the problems that we need to solve:

We need to gain access to the Device Private Key from a genuine Chromecast device. Failing that, we need to coax a device to sign a Peer Certificate for us.
We need to find the Device Certificate that corresponds to the Device Private Key.
We need to find the ICA Certificate that was used to sign the Device Certificate.

Given the statements above - in particular the restricted access to the Device Private Key - this may appear impossible.

Prior Work

While I was initially investigating the Google Cast protocol, I discovered other Chromecast Receiver implementations that did not appear to be officially licensed from Google. Somehow they have solved this problem, so I knew it would be possible to bypass or spoof the Device Authentication mechanism.

My investigation focused on two products, AirServer and Reflector. Using Wireshark and an executable reverse engineering tool called Hopper for macOS, I dug into how both products work internally. Ultimately, it was Reflector that allowed me figure out how to implement a viable Device Authentication mechanism.

Certificate Manifest

The Chromecast receiver functionality in Reflector depends upon something that I’ve been calling a Certificate Manifest. This is a misleading name, as it contains more than just certificates, but I haven’t been able to come up with anything better.

The Certificate Manifest provides all of the information contained in the blue boxes in the Device Authentication diagram shown earlier:

Peer Private Key and Peer Certificate - Used to secure the initial TLS connection. Generated periodically by a Chromecast Receiver.
ICA Certificate - This represents the Intermediate Certificate Authority (ICA) that issues device certificates. This is used for traditional PKI certificate chain verification. Each generation / product line of Chromecast has its own ICA.
Device Certificate - This is a per-device certificate, corresponding to the private key embedded in the device’s secure enclave.
Peer Certificate Signature - Used to ensure that the device is a genuine or licensed Chromecast Receiver. This is signed using the private key embedded within any genuine Chromecast device.

All together, that is five separate items that a receiver needs to pass Device Authentication.

TLS Server

The Peer Private Key and Peer Certificate are the most basic requirements here. These are needed to start a TLS server that can accept connections on port 8009.

Once we are able to accept connections on port 8009, the next problem we need to solve is the Device Authentication Challenge. This is where a Chromecast sender expects us to prove that our receiver is a genuine, or licensed, Chromecast device. To do this, we need to figure out how to sign the Peer Certificate using the Private Key that is embedded in a real Chromecast device.

And in order to do this, we’ll need to gain root access to the Chromecast…

Root Access

A little research was enough to discover that there is a way to gain root access to a Chromecast. The catch is that this was only possible with a first-generation Chromecast, and even then, only one that had not been updated to the latest firmware (i.e. never connected to the internet).

Amazingly, I was able to find not one, but two mint condition first-gen Chromecast units online, and quickly ordered these for safe keeping:

I apologise for hoarding these. For this project, it was crucial that I had a backup device, in case the first was bricked while doing research. Of course, I have been very careful to avoid connecting either device to the internet with the stock firmware. Even a few minutes of internet connectivity would be enough for the device to download a firmware update that closes the vulnerability allowing root access.

The first generation Chromecast was not just physically vulnerable, but at the time, the Chromecast firmware was also vulnerable to an attack in which the Chromecast could be booted using an external USB device.

Eureka-ROM Firmware

Eureka-ROM is a third-party Chromecast ROM, maintained by Team Eureka. It includes features such as ADB, SSH and Telnet access, and DNS customization. Other parts of the OS remain unchanged. The ROM is playfully named after the original codename for the Chromecast project: Eureka.

Eureka-ROM also includes and a web-based control panel, which can be useful for debugging rooted devices:

Crucially, we’ll be relying on ADB and SSH access in order to run our own binaries on a rooted Chromecast.

Flashing Process

The actual process of rooting the device is interesting. The process requires a Teensy++, a simple USB-based micro-controller development board that we will connect to the Chromecast via USB. The Teensy++ allows us to boot the Chromecast into a mode that allows a new firmware image to be written to the device.

The first thing we need to do is program the Teensy++ with firmware (called HubCap) that allows it to communicate with the Chromecast. This is done by connecting the Teensy++ to a PC and using Teensy Loader to program the micro-controller. We also need to prepare a USB thumb drive with the firmware (Eureka-ROM) that we intend to install on the Chromecast. This is not a regularly formatted USB thumb drive. Instead, we need to use a disk imaging tool to write the firmware image to the drive.

With all this ready, we can connect the Teensy++ to the Chromecast, using an externally-powered USB OTG cable:

Once the Chromecast has been booted into a mode that allows us to upload new firmware, we replace the Teensy++ with the USB thumb drive containing our custom firmware. Updating the firmware on the Chromecast takes up to 10 minutes, so when I was testing this, I would just set a timer and go do something else in the meantime. If everything has worked as intended, when we next boot the Chromecast, we’ll see this lovely splash screen:

Gotchas

It’s important to point out that, at time of writing, these steps will render your Chromecast useless, unless you go through the steps to re-flash it with official firmware.

I found that versions 17977.001 and 44433.001 of Eureka-ROM are both too old to be compatible with modern versions of Google Chrome. My attempts to upgrade to later versions were entirely unsuccessful. As such, it’s best to think of any Chromecast used for signature generation as a donor device.

For reasons explained below, I settled for version 17977.001 of the Eureka-ROM firmware.

Signature Generation

We’re almost there!

As mentioned above, a Chromecast Receiver is expected to generate a new TLS private key and self-signed certificate at least once every 48 hours. I generally refer to these as the Peer Private Key and Peer Certificate.

When a Device Authentication Challenge is received, we need to calculate the SHA1 hash of the Peer Certificate (this will be X bytes long) and sign this using Device Private Key.

Chromecast Tools

When first researching this project, I found a GitHub repo called chromecast-widevine-tools. This contains source code for tools that can be run on a rooted Chromecast in order to use the cryptographic functionality of the device, including signing a hash using the private key. I have no idea how the original author came up with any of this, but remarkably, it works.

Compiling these tools was a little tricky, because I had to obtain an appropriate toolchain, and set up the necessary development packages on a Linux system. At the conclusion of this project, I created a fork of this repo, chromecast-tools, which contains just the code needed to sign data using the private key. This makes it sufficient for Device Authentication.

My fork also includes a Dockerfile to make compiling a bit easier:

git clone git@github.com:tristanpenman/chromecast-tools.git
cd chromecast-tools
docker build . -t chromecast-tools
docker run -v `pwd`:/workspace --rm -it chromecast-tools /bin/bash
make

This produces a binary called gtv-ca-sign, which can be run on the Chromecast device.

Aside: Toolchain Compatibility

At this point, it’s worth mentioning some issues that I had with toolchain compatibility.

In the Gotchas section above, I mentioned that I tried versions 17977.001 and 44433.001 of the Eureka-ROM firmware and ultimately settled for 17977.001. The reason for this is toolchain compatibility.

A toolchain contains all the libraries, headers, and other supporting files required to compile programs for a specific architecture and operating system. When cross-compiling (as we are here) it is crucial that the toolchain is compatible with the libraries and OS installed on the target device.

What I found is that between versions 17977.001 and 44433.001, some kind of breaking changes had occurred. The result is that the gtv-ca-sign binary would be compiled successfully, but running it on the device would fail with an error like this:

gtv-ca-sign: not found

I chose not to investigate this further, and simply revert to the working firmware.

Mining Certificates

This is the final piece of the puzzle.

The first step in successfully passing Device Authentication is to start a TLS server using a Peer Private Key and self-signed Peer Certificate that have been generated within the last 24 hours, and are valid for no longer than 48 hours. A Chromecast Sender will immediately close the connection if these conditions are not met!

We also need to generate a cryptographic hash of the Peer Certificate, which we will sign using the Private Key belonging to our rooted Chromecast.

I wrote a little Python script to generate an appropriate Peer Private Key and Peer Certificate. This script also calculates the SHA256 hash of the Peer Certificate:

from OpenSSL import crypto
from datetime import date, datetime, time, timedelta, timezone
import hashlib
import os
import sys

today = date.today()

def get_iso_timestamp(day_offset):
  timestamp_string = datetime.combine(today + timedelta(days = day_offset),
                                      time(0, 0, 0),
                                      tzinfo=timezone.utc).strftime('%Y%m%d%H%M%S')
  return str.encode(timestamp_string + 'Z')

# Generate a key
k = crypto.PKey()
k.generate_key(crypto.TYPE_RSA, 2048)

# Start building our certificate
cert = crypto.X509()
cert.get_subject().CN = 'FakeChromecast'
cert.set_serial_number(12345678)
cert.set_pubkey(k)
cert.set_issuer(cert.get_subject()) # self-signed

# Set the validity period
cert.set_notBefore(get_iso_timestamp(-1)) # 24 hours ago
cert.set_notAfter(get_iso_timestamp(1)) # 24 hours from now

# Sign the cert
cert.sign(k, 'sha1')

# Write the cert to a file
f = open(os.path.join(target, 'peer.crt'), 'wb')
f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert))
f.close()

# Write the key to a file
f = open(os.path.join(target, 'peer.key'), 'wb')
f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, k))
f.close()

# Print out a SHA256 hash of the Peer Certificate
der = crypto.dump_certificate(crypto.FILETYPE_ASN1, cert)
sha = hashlib.sha256(der).hexdigest()
print(sha)

That last bit (printing the hash) is important, because we need to sign the SHA256 hash of the Peer Certificate using the private key embedded in our rooted Chromecast. We do this by passing the hash to gtv-ca-sign (replacing <HASH> with the hash that was printed by the script):

adb shell /tmp/gtv-ca-sign --hash <HASH>

This will produce the signature that we need when responding to a Device Authentication challenge!

Conclusion

When I first started this project (several years ago) it was intended to be a small proof-of-concept. In the years since, it became an interesting personal project that challenged me to learn Go. While my implementation could only ever be considered a toy, it helped me to internalise important concepts relating to PKI and the security of physical devices.

This has become relevant once again, courtesy of the recent Chromecast outage. Crucially, doing what we can to spread knowledge about how these security mechanisms work will help ensure that future devices are designed without such vulnerabilities.

References

Open source projects relevant to this post:

Go-Cast - My own Chromecast Receiver implementation, that I wrote while learning Go. It implements Device Authentication, basic app discovery, and rudamentary screen mirroring using VP8. At time of writing, the RTP implementation is still very buggy…
Chromecast Tools - My fork of chromecast-widevine-tools that contains just the code needed for generating certificates and signing Peer Certificates.
Open Screen Library - Google’s Open Screen Library, which includes implementations for Google Cast Senders and Receivers. This is more complete than my Go implementation, and would be an excellent reference for those wanting to implement Google Cast themselves.

Related projects found during the proof-of-concept phase:

node-castv2 - An implementation of the Google Cast protocol in JavaScript
node-castv2-client - A corresponding client implementation written in JavaScript

Interesting reading:

The Chromecast 2’s device authentication certificate has expired
Google has a fix for your broken Chromecast V2 unless you factory reset
Forensic analysis of Chromecast and Miracast devices - Peter van Bolhuis, Cedric Van Bockhaven

Roll Your Own Ruby Type Checking: Part 3

Sat, 20 May 2023 00:00:00 +0000

In parts one and two of this series, we built a basic runtime type checker for Ruby. The design of our type checker is heavily inspired by Sorbet. So much so that we run up against some of the same limitations as Sorbet. This post focuses on one of those limitations - the inability to perform type checks on parameters with default values. To better understand this, we’ll look at several examples in both Ruby and Python.

Ruby Parameter Types

Recall from part two, the six parameter types that our type checker is currently designed to support. We used reflection to derive this list:

req - required positional
opt - optional positional (has default value)
rest - additional positional arguments
keyreq - required keyword parameter
key - optional keyword parameter
keyrest - additional keyword parameters

The two that concern us in this post are opt and key, as both of these represent parameters that have a default value.

So what’s the actual issue here? Well, it turns out that neither Sorbet, nor our type checker, are able to check the types of those default values!

Breaking Sorbet

I originally wanted to title this post Breaking Sorbet and to publish it independently of this series. But it wouldn’t be fair to call out Sorbet for this, as it is a limitation that stems from decisions made during the design of Ruby’s optional parameter types. In fact, I only discovered this issue while testing my own type checker under different scenarios.

To see how this affects type checking in Sorbet, check out the following example:

require 'sorbet-runtime'

class Main
  extend T::Sig

  sig do
    params(kw: String).void
  end
  def self.main(kw = Time.now)
    puts "kw: #{kw}"
    puts "kw.is_a?(String): #{kw.is_a?(String)}"
  end
end

Here we’re using Time.now as the default value for the main method, and this will be evaluated every time the method is called.

First we’ll try calling the main method with an invalid argument:

begin
  Main.main(123)
rescue StandardError => e
  puts "Error: #{e}"
end

This will fail as expected:

Error: Parameter 'kw': Expected type String, got type Integer with value 123
Caller: ./part-3.rb:36
Definition: ./part-3.rb:29

If we call it without an argument, it will use the default value (Time.now) that we’ve defined:

loop do
  Main.main
  sleep 1
end

Somehow, this works:

kw: 2023-05-13 09:49:07 +1000
kw.is_a?(String): false
kw: 2023-05-13 09:49:08 +1000
kw.is_a?(String): false
kw: 2023-05-13 09:49:09 +1000
kw.is_a?(String): false

We can see that Sorbet’s type check is satisfied, even though we can clearly see that the value is not a Ruby String.

But why?

So why is the type check succeeding here?

Recall that our type checker works by wrapping a method. The wrapper knows how to check the types of its arguments, before passing those as arguments to the original method. This works fine when there are arguments to check. But in the case of default values, there arguments are absent. And this means that the Ruby interpreter evaluates the expression for the default value as the original method is called. The only way we would be able to inspect that value would be to inject code into the method, which is far beyond the scope of our type checker.

In a nutshell, default values are evaluated at the time that a method is invoked.

The Python Way

Interestingly, this is not necessarily how other languages handle default values… In Python, for example, default values are evaluated at the time that a method is defined. We can run a little experiment (in Python of course) to see how this behaves:

def arr():
  return [1,2,3]

def test(a = arr()):
  a.append(4)
  print(a)

test()
test([8,9,10])
test()

You might expect this to produce the following:

[1, 2, 3, 4]
[8, 9, 10, 4]
[1, 2, 3, 4]

But instead we get this:

[1, 2, 3, 4]
[8, 9, 10, 4]
[1, 2, 3, 4, 4]

This shows us that:

The array has been passed to the method body by-reference.
The arr() method is only called once, at the time that the test method was defined.

This can lead to bugs that are very confusing and hard to track down, and it is notorious for tripping up novice Python programmers!

Solutions?

Is there anything we can do about this? Well, maybe…

It may be possible with some significant changes to the Ruby interpreter, such as the addition of more powerful reflection and metaprogramming functionality. This would allow purely Ruby-based type checkers greater access to the method invocation process. This could have applications for performance monitoring and other low-level problems.

Another approach that I would like to explore is the use of runtime code generation - effectively reading the bytecode that is generated for a method, and to modify it such that the type checking occurs inside the method body. However this lies far outside the scope of this series, so don’t expect to see that for while!

Roll Your Own Ruby Type Checking: Part 2

Sat, 13 May 2023 00:00:00 +0000

In my last post on Ruby Type Checking, we looked at how metaprogramming could be used to implement an annotation-style syntax in Ruby. We then leveraged these techniques to build a simple runtime type checker, heaviliy inspired by Sorbet. Unfortunately, our type checker is woefully incomplete. In this post, we take another step towards making it useful.

The code for this post can be found in my ruby-type-checking repo on GitHub.

Recap

The examples in the first post were based on a toy class called Repeater. This class includes one method, repeat, whose sole purpose is to repeat a string (str) a given number times (count), with each copy of the string separated by another string (separator). We made this more interesting by deciding that str and count should be positional parameters, while separator would be a keyword parameter.

By the end of the first post, our final type checked example looked like this:

class Repeater4
  extend Types

  typedef do
    params(
      String, Numeric, separator: String
    ).returns(
      String
    )
  end
  def repeat(str, count, separator: '')
    Array.new(str, count).join(separator)
  end
end

Limitations

Our type checker was complete enough to support the Repeater example. However, it is still limited in various ways. In particular, it is unable to represent the full range of parameter types supported by Ruby. At a high level, these include:

Rest Parameters - Rest Parameters allow an arbitrary number of arguments to be passed to a method. e.g. the method def my_func(tag, *args) requires one argument, followed by zero or more additional arguments. Our type checker does not support this.
Optional Parameters - These are positional parameters or keyword parameters with default values.
Block Parameters - Methods can accept a block as an argument, and can call yield to pass control to that block. This is commonly used for methods like each or map on arrays.

Our type system is also quite simplistic, and it isn’t difficult to imagine a couple of ways in which it could be improved:

Compound Types - It would be nice to be able to represent compound types. These are types that come from combining other types, e.g. OneOf(String, Numeric).
Array and Hash Types - It is also common to pass arrays and hashes as method arguments, and we would like some way to represent this, e.g. Array(String) or Hash(key1: String, key2: Numeric).

Addressing all of these limitation would be too much for one post, so we’ll focus on Rest Parameters and Optional Parameters for now. This alone will bring us much closer to having a useful type checker.

Adding support for Blocks Parameters, or Compound, Array and Hash Types will be the topic of a future post.

Rest Parameters

Lets begin with Rest Parameters. If we revisit the code from our first type checker, we can see that we validate positional arguments by iterating over an array of expected argument types:

def check_types(args, kwargs, arg_types, kwarg_types)
  # check positional arguments
  arg_types.each_with_index do |type, idx|
    raise "Invalid type for arg #{idx}; expected: #{arg_types[idx]}" \
      unless args[idx].is_a? type
  end

  # check keyword arguments
  kwarg_types.each do |key, type|
    raise "Invalid type for kwarg '#{key}`; expected #{kwarg_types[key]}" \
      unless kwargs[key].is_a? type
  end
end

Keyword arguments are validated using a similar approach, but in that case we’re iterating over the keys in a Hash.

This implementation assumes that a method will always accept a fixed number of positional arguments, followed by zero or more keyword arguments, like our repeat method:

def repeat(str, count, separator: '')
  Array.new(str, count).join(separator)
end

Unfortunately, this isn’t the complete picture. We may have a method like send_messages below, that accepts a single positional argument, followed by a variable number of positional arguments:

def send_messages(recipient_id, *messages)
  messages.each do |message|
    send_message(recipient_id, message)
  end
end

The asterisk (splat operator) means that, within our method body, messages will actually refer to an array containing all additional positional arguments.

How does Sorbet do it?

Clearly we’ll need to take a different approach, that allows us to annotate *messages with a type, even though it may represent zero or more positional arguments.

The way Sorbet approaches this is to assign types to parameter names, even those that are positional. Therefore, the type for a variable-length argument list such as *messages is a type that must be satisfied by all elements in the array. The Sorbet solution looks like this:

sig do
  params(
    # Integer describes a required positional argument
    recipient_id: Integer,
    # String describes every single value of the `messages` array,
    # since `messages` uses the splat operator
    messages: String
  ).void
end
def send_messages(recipient_id, *messages)
  messages.each do |message|
    send_message(recipient, message)
  end
end

We’ll need to make some changes to our type checker to support this.

Optional Parameters

The other case we are going to address is optional parameters. Optional parameters are those that have default values, and may be omitted. There are two kinds that we need to cater for:

Optional positional parameters - These must come after required positional parameters, but before additional/rest parameters and keyword parameters.
Optional keyword parameters - These must come after positional parameters, but can appear anywhere in a sequence of keyword parameters.

This makes things more complicated. At this stage, it’s worth digging into the different kinds of method parameters that Ruby supports. One way we can do this, is to use reflection to inspect the parameters for a method.

Parameter Types

Here is a simple example, which includes six different parameter types (blocks are excluded). This method uses reflection to print out the type for each of its parameters:

def test1(a, b = 2, *c, d:, e: 6, **f)
  method(__method__).parameters.each do |param|
    puts param.to_s
  end
  nil
end

When we call it:

test1(1, 2, 3, 4, d: 5, e: 6, f: 7, g: 8)

we’ll get a list of parameters and what kind they are:

[:req, :a]
[:opt, :b]
[:rest, :c]
[:keyreq, :d]
[:key, :e]
[:keyrest, :f]

This gives us six kinds:

req - required positional
opt - optional positional (has default value)
rest - additional positional arguments
keyreq - required keyword parameter
key - optional keyword parameter
keyrest - additional keyword parameters

The weird naming here is a reflection of Ruby’s history, with different keyword parameter types being introduced at different times.

Binding

Now we can go a step further, and use reflection to see how arguments are bound to parameter names:

def test2(a, b = 2, *c, d:, e: 6, **f)
  method(__method__).parameters.each do |param|
    puts [
      # parameter name
      param[1].to_s,
      # argument
      binding.local_variable_get(param[1].to_s)
    ].to_s
  end
  nil
end

Let’s see what happens when we omit arguments for any optional parameters:

test2(1, d: 5, f: 7, g: 8)

Running this, we get:

["a", 1]
["b", 2]
["c", []]
["d", 5]
["e", 6]
["f", {:f=>7, :g=>8}]

We can see that b and e have been bound to their default values. We can also see that the rest parameter c is bound to an empty array, as there were no extra positional arguments.

New Type Checker

We can put all of this together to write a new version of our type checker. Here is the core of the Types module:

module Types
  # we accept just a hash of key-values, mapping parameters to argument types
  def params(**arg_types)
    @arg_types = arg_types
    self
  end

  def returns(ret_type)
    @ret_type = ret_type
    self
  end

  def typedef
    yield
  end

  def method_added(name)
    # short-circuit, to avoid infinite loop
    return unless @arg_types || @ret_type

    # reset hook context, but store current values
    arg_types, ret_type = @arg_types, @ret_type
    @arg_types, @ret_type = nil, nil

    # capture the original method
    meth = instance_method(name)
    params = meth.parameters

    # wrap the original method with type checks
    define_method(name) do |*args, **kwargs, &block|
      Helpers::check_positional_args(args, arg_types, params)
      Helpers::check_keyword_args(kwargs, arg_types, params)
      ret = meth.bind(self).call(*args, **kwargs, &block)
      Helpers::check_return_value(ret, ret_type) if ret_type
      ret
    end
  end
end

To make this more maintainable, the actual type checking code has been moved into a Helpers module. It’s here that we implement our new name-based type checker:

module Types
  module Helpers
    class << self
      def check_positional_args(args, arg_types, params)
        arg_type = nil
        args.each_with_index do |arg, idx|
          param = params[idx]
          if param && [:req, :opt, :rest].include?(param[0])
            param_type = param[0]
            param_name = param[1]

            # Updated as long as their are positional parameter names to
            # consume. Once there aren't any more to consume, we must be
            # checking additional arguments, and can keep using whatever
            # the last type was for that.
            arg_type = arg_types[param_name] unless arg_types[param_name].nil?

            # Happens if there are positional arguments without a corresponding
            # type for the rest parameter
            raise "Missing type for #{param[0]} parameter `#{param_name}`" \
              if arg_type.nil?
          end

          raise "Invalid arg at position #{idx}, expected #{arg_type}" \
            unless arg.is_a?(arg_type)
        end
      end

      def check_keyword_args(_kwargs, _arg_types, _params)
        # TODO: not implemented yet
      end

      def check_return_value(ret, ret_type)
        raise "Invalid return type, expected #{ret_type.name}" \
          unless ret.is_a? ret_type
      end
    end
  end
end

Keyword Parameters

Finally, we can tackle validation of keyword parameters, using our new approach. This is slightly trickier than it may sound, as we would like to detect when a method is called with additional keyword arguments that are either unexpected, or of an incorrect type.

To do this, we make a copy of the keyword argument hash supplied to the method, and we remove arguments from this hash as they are determined to be correct. Anything remaining at the end must be an additional keyword argument, and should be validated using the additional keyword argument type (if specified).

We can re-open our Helpers module, to implement this:

module Types
  module Helpers
    class << self
      def check_keyword_args(kwargs, arg_types, params)
        # don't modify the original hash
        kwargs = kwargs.clone

        arg_type = nil
        params.each do |param|
          param_type = param[0]
          param_name = param[1]
          arg_type = arg_types[param_name] unless arg_types[param_name].nil?

          if param_type == :keyrest
            # only have keyrest params left, so we can break out
            break
          elsif param_type == :keyreq
            raise "Invalid value for required kw param `#{param_name}`; expected #{arg_type}" unless \
              kwargs[param_name].is_a?(arg_type)
          elsif param_type == :key
            raise "Invalid value for optional kw param `#{param_name}`; expected #{arg_type}" unless \
              !kwargs.include?(param_name) || kwargs[param_name].is_a?(arg_type)
          else
            raise "Unexpected param type: #{param_type}"
          end

          # make sure we can detect extra kw params when they're not expected
          arg_type = nil
          kwargs.delete(param_name)
        end

        raise 'Unexpected extra kw params' \
          if kwargs.keys.length > 0 && arg_type.nil?

        kwargs.keys.each do |kwarg|
          raise "Invalid value for extra kw param `#{kwarg}`; expected #{arg_type}" unless \
            kwargs[kwarg].is_a?(arg_type)
        end
      end
    end
  end
end

Closing Thoughts

What I hope you take away from this post is that very changes are required to make our type checker work with a wide range of Ruby parameter types. We saw how easy it is to use reflection to ticker with Ruby’s parameter binding logic, and it was by conducting these experiments that we discovered the tricks that would allow us to implement our type checker.

Stay tuned for the next post in this series, in which we’ll explore some of the limits of runtime type checking in Ruby…

Porting Asteroids to the Nintendo 64

Sat, 31 Dec 2022 00:00:00 +0000

It was during Melbourne’s record-breaking 2020 lockdown that I rediscovered my interest in retro console development. While I had previously dabbled in homebrew for consoles from the 16-bit era, it was now the later generation consoles, such as the Nintendo 64 and SEGA Dreamcast, that caught my interest. These consoles represented the earliest days of the 3D era, each with their own quirks.

This post details my experiences porting Asteroids to the Nintendo 64, and includes some pointers on how you can get started with N64 Homebrew.

Introduction
Inspiration
N64 Architecture
Emulation
Dev Enviornment
Real Hardware
N64 Operating System
NuSystem
Display Lists
Sound Tools
Closing Thoughts

Introduction

It turns out that Nintendo 64 homebrew is now very accessible, with both of official and open source SDKs available. After discovering this, I tried porting my Asteroids clone, with the goal of having it run on real hardware. All of the code for this project can be found on GitHub, and a working N64 ROM can be found here.

This is what it looks like in action:

Inspiration

I first became aware that there was an active N64 homebrew scene after accidentally discovering James Friend’s goose64, a de-make of Untitled Goose Game. His code made use of the official N64 SDK and the Windows-based development kit that many game developers would have used back in the 90s.

I thought the demake project was a cool idea, and the fact that it was open source was extremely valuable in getting me started on my own projects. While exploring James’ GitHub profile, I found that he had a number of repos dedicated to N64 Homebrew. In particular, the n64-sdk-demo repo is a good starting point for budding N64 developers.

Another good starting point is the n64 and N64Homebrew subreddits, and Discord:

Lets move on to take a look at the Nintendo 64 architecture…

N64 Architecture

The Nintendo 64 is built around two main processors. The main CPU is based on the MIPS R4300i, clocked at 93.75 MHz. The CPU is complemented by the RCP (or Reality Coprocessor) which is used for graphics and audio.

The system includes two Rambus RAM chips (RDRAM), providing a total of 4.5MB of system memory (0.5MB reserved for the RCP). System memory can be expanded to 9MB (8MB usable) with the N64 Expansion Pak. The N64 features a Unified Memory Architecture, meaning that the same memory is used by both the CPU and RCP.

CPU

The R4300i is a 64-bit CPU, and is considered the successor to the 32-bit R3000, the CPU that was used in the original Sony PlayStation.

RCP

The RCP was designed by SGI, specifically for the Nintendo 64. It consists of two main cores. One is the Reality Signal Processor, or RSP. The RSP is a vector processor based on the MIPS R4000 design and clocked at 62.5 MHz. The other is the Reality Display Processor, or RDP. While the RSP is primarily used for vector processing and mixing audio, the RDP is responsible for rasterisation.

The RCP is also used to handle joystick/controller inputs, and serves as a DMA controller for memory access. The importance of this chip is highlighted by its position on the N64 motherboard, where it occupies an even more central location than the CPU.

If you want to learn more about the N64 Architecture, then I can highly recommend Rodrigo Copetti’s write-up:

Nintendo 64 Architecture - A practical analysis by Rodrigo Copetti

His website is an absolute treasure trove of information about retro game consoles!

Graphics

One misconception that I held, going into this project, was that the N64 graphics APIs would closely resemble OpenGL (since both APIs were developed by SGI). In practice, there were a number of key differences that needed to be taken into account.

The N64 graphics APIs are based on the concept of a Display List, which is submitted to the RSP to be processed, and then rasterised by the RDP. Display Lists are limited in size, and must carry the geometry that is to be drawn. This is quite different to the OpenGL 1.x approach, where the GL implementation can be thought of as a giant state machine, relying on direct calls to glLine, etc, or otherwise referencing vertex arrays in memory.

Display Lists are discussed in more detail later in this post.

As a result, my existing graphics code would have to be re-written to work with the N64 graphics pipeline. The biggest challenge I ran into was porting my rudimentary font system. I quickly ran into resource limitations when naively copying font data into Display Lists, and had to put some effort into doing this more efficiently.

Here is an example of some graphics code from my n64-homebrew repo on GitHub:

// A static array of model vertex data.
//
// This includes:
//  - position (x,y,z)
//  - texture U,V coords (called S,T in the SDK)
//  - vertex color values in r,g,b,a form.
//
// As this data will be read by the RCP via direct memory access,
// which is required to be 16-byte aligned, it's a good idea to
// annotate it with the GCC attribute `__attribute__((aligned (16)))`,
// to force it to be 16-byte aligned.
//
Vtx vertices[] __attribute__((aligned (16))) = {
  //  x,   y,  z, flag, S, T,    r,    g,    b,    a
  // front
  { -64,  64, 64,    0, 0, 0, 0x00, 0xff, 0x00, 0xff  },
  {  64,  64, 64,    0, 0, 0, 0x00, 0x00, 0x00, 0xff  },
  {  64, -64, 64,    0, 0, 0, 0x00, 0x00, 0xff, 0xff  },
  { -64, -64, 64,    0, 0, 0, 0xff, 0x00, 0x00, 0xff  },
  // back
  { -64,  64, -64,   0, 0, 0, 0x00, 0xff, 0x00, 0xff  },
  {  64,  64, -64,   0, 0, 0, 0x00, 0x00, 0x00, 0xff  },
  {  64, -64, -64,   0, 0, 0, 0x00, 0x00, 0xff, 0xff  },
  { -64, -64, -64,   0, 0, 0, 0xff, 0x00, 0x00, 0xff  },
};

void stage1_draw_box()
{
    // load vertex data for the triangles
    gSPVertex(gfx_list_ptr++, &(vertices[0]), 8, 0);

    // depending on which graphical features, the RDP might need to
    // spend 1 or 2 cycles to render a primitive, and we need to
    // tell it which to do
    gDPSetCycleType(gfx_list_ptr++, G_CYC_1CYCLE);

    // use antialiasing, rendering an opaque surface
    gDPSetRenderMode(gfx_list_ptr++,
        G_RM_AA_ZB_OPA_SURF, G_RM_AA_ZB_OPA_SURF2);

    // reset any rendering flags set when drawing the previous primitive
    gSPClearGeometryMode(gfx_list_ptr++, 0xFFFFFFFF);

    // enable smooth (gourad) shading and z-buffering
    gSPSetGeometryMode(gfx_list_ptr++,
        G_SHADE | G_SHADING_SMOOTH | G_ZBUFFER);

    // actually draw the triangles, using the specified vertices (front)
    gSP2Triangles(gfx_list_ptr++, 0, 1, 2, 0, 0, 2, 3, 0);

    // back
    gSP2Triangles(gfx_list_ptr++, 5, 4, 7, 5, 5, 7, 6, 5);

    // right
    gSP2Triangles(gfx_list_ptr++, 6, 5, 1, 6, 6, 1, 2, 6);

    // left
    gSP2Triangles(gfx_list_ptr++, 0, 7, 3, 0, 0, 4, 7, 0);

    // Mark that we've finished sending commands for this particular
    // primitive. This is needed to prevent race conditions inside
    // the rendering hardware in the case that subsequent commands
    // change rendering settings.
    gDPPipeSync(gfx_list_ptr++);
}

This code is part of an example that renders spinning box:

You can see that for the gDP and gSP API calls, the first argument is always a pointer into a Display List. This is post-incremented after every call, so that gfx_list_ptr points to the next empty slot in the Display List.

Audio

Audio is another case where the Nintendo 64 is a bit unusual. Whereas most home consoles would have a dedicated sound chip, or audio processing pipeline, the Nintendo 64 uses the RSP for processing audio. Similar to how Display Lists are used for graphics processing, the RSP uses Audio Command Lists for audio.

Audio can be thought of as a layered architecture, with the SDK providing ‘Players’ that are managed by the CPU. The output of these players is fed to a Synthesis driver that serialises this data in a format understood by the audio microcode running on the RSP.

Input

Cross-platform input handling required significant refactoring. Whereas the original code was event based, relying on the OS to publish input events on an event loop, the N64 expects programmers to poll for the current state of input devices.

To make this consistent across the N64 and desktop versions of the code, I chose to implement an input mapping scheme. This made it possible to map multiple platform-specific inputs to a single input ID/handle. The N64 input handling code used a brute-force approach, where all relevant inputs are checked once per frame. The game code was already based on a fixed time step, so this worked well, and the overhead did not cause performance issues.

Storage (Controller Pak)

There are two ways we can manage high scores. The first would be to store them in save stage on the cartridge, relying on SRAM built into the cartridge. The second is to store them using a Controller Pak, which would allow the save data to be transferred between consoles/cartridges. I opted for the Controller Pak approach.

This proved to be quite challenging, as the APIs for accessing the Controller Pak require certain initialisation steps to be taken before attempting to read/write from the Controller Pak. Failing to do so leads to unpredictable behaviour, sometimes even corrupting data on the Controller Pak. I spent more time on this than I would have liked, relying on printf debugging via the EverDrive serial connection to figure out what was going wrong.

As you can see from the following code snippet, just initialising the Controller Pak is quite intricate:

bool open_controller_pak()
{
    s32 ret;

    nuContPakCodeSet(COMPANY_CODE, GAME_CODE);
    nuContQueryRead();
    nuContDataReadStart();
    nuContDataReadWait();

    if (nuContStatus[0].status & CONT_CARD_ON) {
        debug_printf(" - nuContStatus says card is on\n");
    } else {
        debug_printf(" - nuContStatus says card is off\n");
        return false;
    }

    nuContPakOpen(&pak_file, 0);
    if (pak_file.type == NU_CONT_PAK_TYPE_PAK) {
        debug_printf(" - nuContPakOpen says card is controller pak\n");
        return true;
    }

    if (!nuContRmbCheck(0)) {
        debug_printf(" - nuContRmbCheck says card is rumble pak\n");
        return false;
    }

    debug_printf(" - nuContPakOpen error: %d\n", pak_file.error);
    debug_printf(" - nuContPakRepairId attempting repair...\n");
    ret = nuContPakRepairId(&pak_file);
    if (ret != PFS_ERR_NEW_PACK) {
        debug_printf(" - nuContPakRepairId failed: %d\n", ret);
        return false;
    }

    // try again
    nuContPakOpen(&pak_file, 0);
    if (pak_file.type == NU_CONT_PAK_TYPE_PAK) {
        debug_printf(" - repair successful\n");
        return true;
    }

    debug_printf(" - repair not successful\n");
    return false;
}

Haptics (Rumble Pak)

While it may seem crude now, Nintendo’s Rumble Pak was yet another example of the company’s willingness to experiment with new ideas (just look at the shape of the controller itself). Force feedback is now an expected feature in all modern game controllers.

The Rumble Pak requires 2 AAA batteries. Devices that take AAA batteries are typically designed to handle voltages from 0.9 to 1.5V. In practice, this could mean weaker feedback, but would not be a critical issue for the Rumble Pak. When discussing the need for AAA batteries with a friend, we guessed that this was to avoid using/straining the 3.3V that is supplied by the console.

It turns out that 1) it is possible to run the Rumble Pak without batteries and 2) this would be a bad idea when more than two controllers are fitted out with Rumble Paks:

Driving the Rumble Pak was straightforward, using a simple API that allows game code to set the motor speed and duration for feedback events:

void rumble_start(float time, float intensity)
{
    const u16 frequency = 256.0f * intensity;
    const u16 frame = 60.0f * time;

    nuContRmbModeSet(0, NU_CONT_RMB_MODE_ENABLE);
    nuContRmbStart(0, frequency, frame);
}

It’s interesting to note that the N64 controller was ultimately limited by a half-duplex serial interface. This means that instead of separate wires/pins for each button, there is a single data wire that carries all data. The protocol used later become known as the Joybus Protocol.

Emulation

There are many challenges in N64 emulation. It’s not just a matter of emulating the CPU. The RCP and RSP also need to be emulated.

The RCP relies on microcode, which is driven using Display Lists. While this can be emulated on the instruction level, doing so is computationally intensive, even on modern hardware. It is more common to emulate the RCP at the microcode level, using what is known as High-Level Emulation (HLE).

UltraHLE was a great example of High-Level Emulation in action. This emulated achieved what was thought to be impossible - full speed N64 emulation on PC hardware. UltraHLE was designed to intercept C function calls, which were mapped to efficient implementations on the CPU, allowing games such as Ocarina of Time and Super Mario 64 to run efficiently on PC hardware in the late 90s and early 2000s.

While High-Level Emulation works well for many games, especially first-party titles, it falls apart when a ROM uses non-standard microcode, an example being Indiana Jones and the Infernal Machine and Star Wars Episode I: Battle for Naboo.

See this interesting write-up of the effort that went into reverse engineering some custom microcode that Factor 5 developed towards the end of the N64’s commercial lifespan:

The masterpiece graphic microcode behind the Nintendo 64 version of Indiana Jones and the Infernal Machine and Star Wars Episode I: Battle for Naboo

As for more modern emulators, I used a variety for testing:

CEN64 (macOS and Linux) - I had to build this from source, but fortunately that is not difficult to do. Running a ROM will require a BIOS file called pifdata.bin. This is easy enough to find online. On macOS you will also need to install XQuartz.
sixtyforce (macOS) - A nice no-frills emulator that has good speed and compatibility. Registering is optional, but is nice to do if you want to unlock save state support and remove the watermark.
Project64 (Windows) - A solid Windows-based emulator, with strong support for debugging, a highly configurable emulation.

Dev Environment

There are a number of options for setting up an N64 development environment. If you’re comfortable using the original Windows SDK, this can be obtained from ultra64.ca:

https://ultra64.ca/files/software/other/sdks/n64sdk.7z

There are several options for how you actually install and use this SDK. Some people choose to use a Windows XP or Windows 2000 VM. I found that it was more practical to install the SDK via WINE, running on a Linux VM.

Unfortunately, using WINE is no longer viable recent versions of macOS, due to the removal of 32-bit support.

Exactly how you set this up will depend on your desktop environment, etc. So my recommendation is to find a tutorial on YouTube that is suited to your setup.

Modern SDK

If you’re using Linux, then another good option is to use CrashOveride95’s Modern SDK. This Linux-native SDK uses an up-to-date version of GCC, making significantly more sophisticated optimisations available. It also gives you the ability to use more modern C features.

For my development environment, I was able to install the Modern SDK on the same Linux VM that I used for the official SDK running on WINE.

Hybrid Approach

Unfortunately, these two SDKs have different requirements as far as Makefiles and build steps are concerned.

For this project, I wanted to support both SDKs, so I’ve included two Makefiles.

The code has also been written to compile without warnings with both SDKs. This meant adopting a subset of C that works with both SDKs, and is not precisely defind by a single standard. While this is not ideal for code readability, it makes it possible to compare the performance and behaviour of code compiled under both SDKs, which in itself is an interesting exercise.

Real Hardware

Of course, no homebrew project is really complete until it’s running on real hardware.

For this, I use an Everdrive-64 X7 flash cartridge. The console is connected to a regular HDMI monitor using a RetroTINK 2X-MINI, which you can see next to the classic Nintendo Yellow controller:

Not only does this cart support loading ROMs via an SD-card, tt also includes a USB interface, allowing ROMs to be loaded directly from a PC.

To upload ROMs via USB, I used buu342’s UNFLoader. The USB connection supports bi-directional communication, so it is possible to use UNFLoader for USB debugging.

N64 Operating System

Earlier consoles, such as the SNES or SEGA Mega Drive, often expected programmers to work directly with the hardware. Game code would be written in assembly language. While common assembly routines would often be provided as part of a development kit, the programmer was ultimately responsible for all aspects of the system. In some cases, specially customised cartridges would be used to enhance the capabilities of the console (e.g. the Super FX chip used in Star Fox).

The Nintendo 64 offered a more modern approach. The official SDK provides a set of APIs, targeting the C programming language, that together constitute the Nintendo 64 Operating System. The N64 OS was small enough that it could be embedded in a ROM, requiring only a tiny amount of valuable ROM space. The OS provided APIs that address the following concerns:

Concurrency and message passing
Process/thread scheduling
Register and I/O access

As you can see, this does not include graphics or audio!

NuSystem

In practice, the N64 Operating System presents an extremely challenging development environment, so Nintendo also provided a set of higher level APIs for controlling the hardware. These are collectively known as NuSystem, often shortened to NuSys.

As an example of the functionality included in NuSys, these are a some of the functions provided for interacting with a Rumble Pak:

/**
 * Initiates operation of the Rumble Pak
 *
 * contNo   the Controller number
 * freq     the vibration frequency (1 - 256) where 1 is one pulse emitted
            every 256 frames and 256 is continuous movement
 * frame    the duration of the vibration in number of frames
 */
void nuContRmbStart(u32 contNo, u16 freq, u16 frame);

/**
 * Sets the operation mode of the Rumble Pak
 *
 * contNo   the Controller number
 * mode     the Rumble Pak operation mode
 *
 * valid modes:
 *   NU_CONT_RMB_MODE_DISABLE for disabled
 *   NU_CONT_RMB_MODE_ENABLE for enabled
 *   NU_CONT_RMB_MODE_AUTORUN for automatic recognition
 */
void nuContRmbModeSet(u32 contNo, u8 mode);

All functions in NuSys are prefixed according to the subsystem they apply to. As you can see in the code snippet above, functions for working with the Rumble Pak are prefixed with nuContRmb.

It’s worth noting that the N64 OS and NuSys components are not strictly layered here. It’s entirely reasonable for a game to make use of NuSys, while still interacting directly with the N64 OS, when greater control is required.

Display Lists

As mentioned earlier, graphics on the N64 is based on Display Lists (and in the case of audio, Audio Command Lists). This is how tasks are passed from the CPU to the RSP. This allows for a high degree of concurrent processing, whereby the RSP can process graphics and audio tasks while the CPU works on preparing subsequent tasks.

To help illustrate this, I’ve prepared a sequence diagram. This shows how tasks move from the CPU to the RSP, making use of shared memory:

We can also see that the result of RSP processing a graphics task is a stream of low-level draw commands that are sent to the RDP. This is what drives rasterisation.

The use of shared memory is an important consideration here. As mentioned earlier, the N64 has a Unified Memory Architecture (UMA), which means that the same memory is used by the CPU and RCP. The RCP actually serves as the memory controller for the N64, mediating access to system memory.

When writing graphics code for the N64, it is important to recognise that the RSP is actively reading from the same region of memory used to prepare a Display List. Only when the RSP has completed a task can that memory be used for another task.

Sound Tools

There are two ways to prepare audio for inclusion in a ROM - one is to use the Windows-based Sound Tools, and the other is to use command line tools. After converting existing sound samples into a format that could be included in an N64 ROM, I battled with the Windows-based Sound Tools to marshall these into a form that could be embedded in a ROM.

For all of this, I used the Sound Tools running on Linux, via WINE:

Check out that version timestamp!

My experience with the Sound Tools is that they’re incredibly painful to use. Clicking in the wrong location can easier do something completely unexpected. There are command line tools available, and I plan to use those next time I take on a N64 homebrew project.

Closing Thoughts

One of the interesting things to come out of this project is that the structure and quality of the code is better now than it was before I attempted to port it. All of the graphics code needed to be updated, after being modularised as a generic canvas API. The new API can be used to draw simple shapes, including the traditional Asteroids font.

Audio is still lacking in many ways. With the differences in how audio is handled on the N64 (versus SDL_mixer 2.0), it wasn’t clear how I could modularise this in the same way as the graphics/canvas APIs. The main reason for this is that audio files are embedded in the ROM via the linker, whereas the desktop version loads WAV files from disk.

Overall, I’m keen to pick up another N64 homebrew project at some point, next time making use of the 3D capabilities of the console.

Roll Your Own Ruby Type Checking: Part 1

Mon, 26 Dec 2022 00:00:00 +0000

Last Updated: 10 April 2023

Earlier this year, I was talking with a colleague about Sorbet and other type checking features that are available for Ruby. Although we were with familiar with usage of gradual typing in other languages such as TypeScript and Python, neither of us really knew how this would work under the hood.

This post begins with an example from the Sorbet documentation, and builds up the machinery to support a simplistic type checking implementation. While our implementation will not be production-ready like Sorbet, I hope that the exploration will be enlightening.

Later posts will refine the type checking implementation, making it more versatile, and capable of being used with more sophisticated Ruby programs.

Type Checking

But first, what is Type Checking. Or more specifically, what does it mean to check types?

In a dynamically-typed language such as Ruby, the types of variables, parameters, etc, are not known until runtime. That is, they can vary depending on how your program is run. A variable that is a String in some cases, may be Numeric in other cases. The purpose of a type-checker is to detect bugs that may occur due to mis-matching types. An example of this is passing a String as an argument when a number is expected.

Where does Sorbet fit into this?

Sorbet

Sorbet is a powerful type checker, built specifically for Ruby. It is implemented using Ruby language features, unlike other type checkers, which are often based on pre-processors and other tricks. The Sorbet ecosystem also provides plugins for IDE integration and additional development tooling.

Let’s look at a simple example, showing how Sorbet is used… When we visit the Sorbet website we’re greeting with this little example, illustrating just how easy it is to add type checking to both new and existing code:

class Main
  extend T::Sig

  sig do
    params(
      name: String
    ).returns(
      Integer
    )
  end
  def main(name)
    puts "Hello, #{name}!"
    name.length
  end
end

This could also be written more compactly:

class Main
  extend T::Sig

  sig { params(name: String).returns(Integer) }
  def main(name)
    puts "Hello, #{name}!"
    name.length
  end
end

We can see from this example that Sorbet provides type checking via a method called sig, which looks a bit like an annotation that you might see in a language such as Java.

As is common in the Ruby ecosystem, this one line type annotation is built upon some non-trivial Ruby magic. In this post, we’ll attempt to re-invent our own annotation machinery, and use it to build a simple type checker inspired by Sorbet.

Repeater

For the examples in this post, we’ll work with a toy class called Repeater. All this class does is repeat a string str a given number times count, with each copy of the string separated by separator. Each version of this class will be given its own number, e.g. Repeater1, Repeater2.

The original implementation looks like this:

class Repeater
  def repeat(str, count, separator: '') do
    Array.new(count, str).join(separator)
  end
end

This class is very easy to test from IRB, or from a script:

puts Repeater.new.repeat('test', 3, ', ')

In this case, we’ll get this output:

test, test, test

Now let’s say we want to add some functionality to this class, so that we print its positional arguments before the body is executed, and print out the return value before it is returned to the caller.

Here’s the most direct way we could implement this:

class Repeater1
  def repeat(str, count, separator: '')
    puts "before: [\"#{str}\", #{count}]"
    ret = Array.new(count, str).join(separator)
    puts "after: #{ret}"
    ret
  end
end

Now If we were to run the following:

puts Repeater1.new.repeat('test', 3, separator: ', ')

It would produce:

before: ["test", 3]
after: test, test, test
test, test, test

If we now allow ourselves a little bit of indirection, this can be made more generic by lifting the method body into a lambda:

class Repeater2
  def repeat(*args, **kwargs)
    puts "before: #{args}"

    # lambda literal
    fn = -> (str, count, separator: '') do
      Array.new(count, str).join(separator)
    end

    ret = fn.call(*args, **kwargs)
    puts "after: #{ret}"
    ret
  end
end

What is unusual about Repeater2 is that the repeat method no longer uses the original parameter list. Instead, it uses *args and **kwargs. These are called rest parameter, and allow us to capture and print all arguments using puts. The code from the original repeat method is captured within a lambda literal, where we see the original parameter list.

If we were to run the following line of code:

puts Repeater2.new.repeat('test', 3, separator: ', ')

It will produce the same output as Repeater1:

before: ["test", 3]
after: test, test, test
test, test, test

Hooks

Of course, we are working towards an annotation-based style, and would like to simplify Repeater2 so that the code to run before and after the method body is more clearly separated. Ideally something like this:

before ->(*args, **kwargs) do
  puts "before: #{args}"
end
after ->(returns) do
  puts "after: #{returns}"
end
def repeat(str, count, separator: '')
  Array.new(str, count).join(separator)
end

In this snippet, before is a hook that is called with the arguments supplied to the method, but before the body of the method is executed. And after is a hook that is called after the body is executed, with its return value.

If we consider that Ruby class declarations are executed from top to bottom, and existing methods can be called from within the class body, then we see that before and after could be provided by extending a module:

class MyClass
  extend Hooks
  # methods from Hooks are available now
  # ...
end

How does the Hooks module work?

We can see from the snippets above, that we need to define a module that provides before and after methods. We’ll call this module Hooks. The before and after methods are very simple setters that store a lambda in a class variable, so that it can be used later. We’ll call this the hook context for a method.

module Hooks
  def before(tag)
    @before = tag
  end

  def after(tag)
    @after = tag
  end

We can then use Ruby’s powerful meta-programming functionality to enhance method declaration, in such a way that the next method to be defined will execute within our new hook contexr.

We’ll do this by implementing method_added, which is automatically called whenever a method is added to a class:

  def method_added(name)
    # short-circuit, to avoid infinite loop
    return unless @before || @after

    # reset hook context, but store current values
    before = @before
    @before = nil
    after = @after
    @after = nil

    # capture the original method
    meth = instance_method(name)

    # wrap the original method with calls to before and after hooks
    define_method(name) do |*args, **kwargs, &block|
      before.call(*args, **kwargs) if before
      ret = meth.bind(self).call(*args, **kwargs, &block)
      after.call(ret) if after
      ret
    end
  end
end

define_method is used to replace the newly added method with a block that executes our before and after hooks. We have to be careful to clear @before and @after before calling define_method, otherwise this code will keep triggering method_added, and we’ll end up with a stack overflow.

We can now implement Repeater3, extending the Hooks module:

class Repeater3
  extend Hooks

  before -> (*args, **kwargs) do
    puts "before: #{args}"
  end
  after -> (returns) do
    puts "after: #{returns}"
  end
  def repeat(str, count, separator: '')
    Array.new(count, str).join(separator)
  end
end

Lets check that it works as expected, using the same test code as earlier:

puts Repeater3.new.repeat('test', 3, ', ')

This should produce exactly the same output:

before: ["test", 3]
after: test, test, test
test, test, test

Type Checking

The next step is to develop this into a convenient, if somewhat primitive, type checking system.

Our first thought may be to use our Hooks module to write bespoke type checking code like this:

before do |str, count, separator|
  raise "invalid str" unless str.is_a? String
  raise "invalid count" unless count.is_a? Numeric
  raise "invalid separator" \
    unless separator.is_a? String || separator.nil?
end
after begin |ret|
  raise "invalid return value" unless ret.is_a? String
end
def repeat(str, count, separator: '')
  Array.new(str, count).join(separator)
end

While this may work, it would be time-consuming, error prone, and not particularly easy to maintain. We would like to streamline this so that we write something more compact, like this:

typedef do
  params(
    String, Numeric, separator: String
  ).returns(
    String
  )
end
def repeat(str, count, separator)
  Array.new(count, str).join(separator)
end

There’s a fair bit going on here. Firstly, we’re no longer using before and after. Instead we’re using a hook called typedef. Within the typedef call, we’re also using methods called params and returns.

This is getting a bit closer to the interface offered by Sorbet:

sig do
  params(
    str: String, count: Integer, separator: String
  ).returns(
    Integer
  )
end
def repeat(str, count, separator)
  Array.new(count, str).join(separator)
end

Method Parameters

One of the interested problems here is that we’re using both positional parameters (str and count) and keyword parameters (separator).

Sorbet uses some additional metaprogramming magic to handle these, however, for the purposes of this experiment, we’ll keep things simple by handling positional and keyword parameters separately.

Types

What we end up with is a Types module that provides two methods, params and returns which can be use to specify the types for a method. These methods both return self, which is what allows them to be chained.

Here is the new Types module in all its glory:

module Types
  def params(*arg_types, **kwarg_types)
    @arg_types, @kwarg_types = arg_types, kwarg_types
    self
  end

  def returns(ret_type)
    @ret_type = ret_type
    self
  end

  def typedef
    # Allows typedef to be passed a block to execute within the context
    # of the Types module
    yield
  end

  def method_added(name)
    # short-circuit, to avoid infinite loop
    return unless @arg_types || @kwarg_types || @ret_type

    # reset hook context, but store current values
    arg_types, kwarg_types, ret_type = @arg_types, @kwarg_types, @ret_type
    @arg_types, @kwarg_types, @ret_type = nil, nil, nil

    # capture the original method
    meth = instance_method(name)

    # wrap the original method with type checks
    define_method(name) do |*args, **kwargs, &block|
      # check positional arguments
      arg_types.each_with_index do |type, idx|
        raise "Invalid type for arg #{idx}; expected: #{arg_types[idx]}" \
          unless args[idx].is_a? type
      end

      # check keyword arguments
      kwarg_types.each do |key, type|
        raise "Invalid type for kwarg '#{key}`; expected #{kwarg_types[key]}" \
          unless kwargs[key].is_a? type
      end

      ret = meth.bind(self).call(*args, &block)

      # check return type
      raise "Invalid return type, expected #{ret.name}" \
        unless ret.is_a? ret_type

      ret
    end
  end
end

Putting it together

Now we can implement the Repeater class using our home-grown type checking system:

class Repeater4
  extend Types

  typedef do
    params(
      String, Numeric, separator: String
    ).returns(
      String
    )
  end
  def repeat(str, count, separator: '')
    Array.new(str, count).join(separator)
  end
end

If we pass valid arguments, as per previous examples, we should see the expected output:

test, test, test

But we replace the count argument with a string, we’ll see our type checking in action:

begin
  puts Repeater4.new.repeat("test", "3", separator: ", ")
rescue StandardError => e
  puts "Error: #{e}"
end

The output should look something like:

Error: Invalid type for arg in pos 1; expected: Numeric

Closing Thoughts

If you take away anything from this post, hopefully it is that Ruby allows you to build annotation-like constructs, using features that are part of the language.

These annotations can be evaluated when classes are loaded, and Ruby’s rich support for metaprogramming means that it is possible to do much more than simply annotate methods.

As for the type checker we constructed in this post, it is currently very simplistic, and not particularly useful. However it lays the groundwork for future posts, in which we’ll develop a more complete type checker.

Code

All of the code for this post can be found in my ruby-type-checking repo on GitHub.

Tristan Penman's Blog

Robotics in Go: Building a Raspberry Pi Rover

Contents

Introduction

The Hardware

Robotics Fundamentals

Robot Operating System

Message Passing

Bare Metal Robotics

Go for Robotics

Concurrency Model

Cross-Compilation

Go Drawbacks

Hardware Interfaces

GPIO

UART

SPI

I²C

Go Libraries

Gobot

Periph

Serial

Rover Kit

MQTT Client

Peripherals

Project Structure

Motor Control

Pulse Width Modulation

Gobot Driver

Periph Driver

Sonar Reader

Sonar Sensors

Voltage Divider

Periph Provider

UART Provider

Web Bridge

WebSocket Server

Gamepad Support

Next Steps

References

Machine Translation at the Edge: Part 2

Contents

Part 1 Recap

MarianNMT

Hugging Face

Hugging Face Transformers

MarianNMT vs MarianMT

Preflight

SentencePiece

Test Script

PyTorch vs ONNX

Converting to ONNX

Marian-ONNX-Converter

Model Files

Converting to RKNN

Quantization

Dynamic Ops

Inference Engine

Initialisation

Tokenization

Encoder

Decoder Loop

Detokenization

Decoding Strategies

Greedy Decoding

Beam Search

Benchmarks

Sample Text

Benchmark Script

Results

Wrap Up

References

Machine Translation at the Edge: Part 1

Contents

Machine Translation

Statistical Models

Aside: Moses

Encoder-Decoder Models

RNNs and LSTMs

Transformers