saveBytes static method

Uint8List saveBytes(Module module, { bool fp16 = false, })

Serialize a module's parameters to a fresh byte buffer.

If fp16 is true, every parameter's data is quantized to IEEE-754 half precision on save (halving file size). The header records "dtype": "F16" per compressed parameter so loadIntoBytes can decode back to fp32 automatically. Old checkpoints without any dtype field continue to load as fp32.

Implementation

static Uint8List saveBytes(Module module, {bool fp16 = false}) {
  final params = module.parameters();
  final header = jsonEncode({
    'version': version,
    'params': [
      for (final p in params)
        {
          'shape': p.shape,
          if (fp16) 'dtype': 'F16',
        },
    ],
    'totalScalars': params.fold<int>(0, (a, p) => a + p.length),
  });
  final headerBytes = utf8.encode(header);

  final bytesPerScalar = fp16 ? 2 : 4;
  final totalScalars = params.fold<int>(0, (a, p) => a + p.length);
  final dataBytes = bytesPerScalar * totalScalars;
  final preamble = 4 + 4 + 4; // magic + version + headerLen
  final out = ByteData(preamble + headerBytes.length + dataBytes);

  // Preamble.
  var off = 0;
  for (int i = 0; i < 4; i++) {
    out.setUint8(off++, magic[i]);
  }
  out.setUint32(off, version, Endian.little);
  off += 4;
  out.setUint32(off, headerBytes.length, Endian.little);
  off += 4;

  // Header JSON.
  for (int i = 0; i < headerBytes.length; i++) {
    out.setUint8(off++, headerBytes[i]);
  }

  // Data blob. Downloading from GPU when needed happens via toList().
  for (final p in params) {
    final vals = p.toList();
    if (fp16) {
      for (int i = 0; i < vals.length; i++) {
        out.setUint16(off, fp32ToFp16Bits(vals[i]), Endian.little);
        off += 2;
      }
    } else {
      for (int i = 0; i < vals.length; i++) {
        out.setFloat32(off, vals[i], Endian.little);
        off += 4;
      }
    }
  }
  return out.buffer.asUint8List();
}