477 lines
18 KiB
Dart
477 lines
18 KiB
Dart
import 'dart:convert';
|
||
import 'dart:io';
|
||
import 'dart:typed_data';
|
||
import '../../l10n/tr.dart';
|
||
|
||
import 'package:path/path.dart' as p;
|
||
|
||
/// ONNX 모델을 메모리가 적은 기기 (Android) 에 맞게 바꾼다. 처음 번역할 때 한 번.
|
||
///
|
||
/// 1. 큰 가중치를 별도 파일 (외부 데이터) 로 옮긴다. 한 파일짜리 모델을 열면 ONNX Runtime 은 파일 내용과
|
||
/// 가중치를 잠깐 함께 들고 있어 모델 크기의 2~3배를 쓴다. 외부 데이터는 파일을 메모리 맵으로 바로 쓰므로
|
||
/// 앱 메모리를 거의 쓰지 않고, Android 도 필요하면 그 메모리를 비울 수 있다.
|
||
/// 2. 정수 lm_head: NLLB 디코더의 마지막 단계는 int8 공유 임베딩 (25.6만 단어 × 1024) 을 매번 전치하고
|
||
/// float 로 풀어 (1GB) 곱한다. 이것을 다른 층과 같은 동적 양자화 정수 곱
|
||
/// (DynamicQuantizeLinear → MatMulInteger → 배율) 으로 바꾸고, 전치한 가중치는 미리 만들어 둔다.
|
||
///
|
||
/// 측정 (NLLB 디코더, 그래프 최적화 끔): 실행 중 최대 1.8GB · 한 단계 800ms → 570MB (대부분 파일 맵) · 60ms.
|
||
/// 바꾼 모델이 원래와 같은 결과를 내는지: tool/check_onnx_external.py
|
||
///
|
||
/// 파일 전체를 메모리에 올리지 않는다: protobuf 머리만 읽으며 큰 raw_data 는 조각조각 .data 파일로 복사하고,
|
||
/// 모델 구조 (수 MB) 만 메모리에서 다시 쓴다.
|
||
class OnnxExternal {
|
||
/// 이보다 큰 가중치만 옮긴다
|
||
static const threshold = 1024;
|
||
|
||
/// .data 안에서 가중치 시작 위치 정렬 (메모리 맵에 유리)
|
||
static const align = 4096;
|
||
|
||
/// 바꾼 모델 경로: encoder_model_quantized.onnx → encoder_model_quantized_jj2.onnx (+ .data).
|
||
/// 바꾸는 방식이 달라지면 숫자를 올린다 (예전 것은 지우고 다시 만든다).
|
||
static String externalPath(String onnx) =>
|
||
p.join(p.dirname(onnx), '${p.basenameWithoutExtension(onnx)}_jj2.onnx');
|
||
|
||
/// 예전 방식으로 바꾼 파일 (정수 lm_head 없음 - 디코더가 1GB 넘게 씀)
|
||
static List<String> _legacy(String onnx) {
|
||
final base = p.join(p.dirname(onnx), '${p.basenameWithoutExtension(onnx)}_ext.onnx');
|
||
return [base, '$base.data'];
|
||
}
|
||
|
||
/// [onnx] 를 [out] (+ [out].data) 로 바꾼다. 중간에 실패하면 만들던 파일을 지운다.
|
||
static void convert(String onnx, String out) {
|
||
final tmp = '$out.tmp', tmpData = '$out.data.tmp';
|
||
final src = File(onnx).openSync();
|
||
final data = File(tmpData).openSync(mode: FileMode.write);
|
||
try {
|
||
final model = _Rewriter(src, data, p.basename('$out.data')).model();
|
||
File(tmp).writeAsBytesSync(model, flush: true);
|
||
} catch (_) {
|
||
src.closeSync();
|
||
data.closeSync();
|
||
for (final f in [tmp, tmpData]) {
|
||
try {
|
||
File(f).deleteSync();
|
||
} catch (_) {}
|
||
}
|
||
rethrow;
|
||
}
|
||
src.closeSync();
|
||
data.closeSync();
|
||
// .data 를 먼저 옮겨 두고 모델 파일을 마지막에 (모델 파일이 있으면 다 된 것)
|
||
File(tmpData).renameSync('$out.data');
|
||
File(tmp).renameSync(out);
|
||
}
|
||
|
||
static bool isReady(String onnx) {
|
||
final ext = externalPath(onnx);
|
||
return File(ext).existsSync() && File('$ext.data').existsSync();
|
||
}
|
||
|
||
/// 바꾼 모델이 없으면 만든다. 바꾼 모델 경로를 돌려준다.
|
||
/// [deleteOriginal]: 다 바꾸면 원래 파일을 지운다 (저장 공간 절약)
|
||
static String ensure(String onnx, {bool deleteOriginal = false}) {
|
||
final ext = externalPath(onnx);
|
||
for (final f in _legacy(onnx)) {
|
||
try {
|
||
File(f).deleteSync();
|
||
} catch (_) {}
|
||
}
|
||
if (isReady(onnx)) return ext;
|
||
convert(onnx, ext);
|
||
if (deleteOriginal) {
|
||
try {
|
||
File(onnx).deleteSync();
|
||
} catch (_) {}
|
||
}
|
||
return ext;
|
||
}
|
||
}
|
||
|
||
/// 정수 lm_head 로 바꿀 곳: Transpose(W) → DequantizeLinear(·, scale, zp) → MatMul(h, ·) → out
|
||
class _LmHead {
|
||
final String weight, transposed, dequantized, scale, zeroPoint, input, output;
|
||
const _LmHead(this.weight, this.transposed, this.dequantized, this.scale, this.zeroPoint, this.input, this.output);
|
||
|
||
String get key => [weight, transposed, dequantized, scale, zeroPoint, input, output].join('|');
|
||
}
|
||
|
||
/// 노드 하나 (분석용)
|
||
class _Node {
|
||
String op = '';
|
||
final inputs = <String>[], outputs = <String>[];
|
||
var perm = <int>[];
|
||
}
|
||
|
||
class _Rewriter {
|
||
final RandomAccessFile src;
|
||
final RandomAccessFile data;
|
||
final String location;
|
||
final _buf = Uint8List(10);
|
||
|
||
/// 정수 lm_head 로 바꿀 곳 (없으면 null)
|
||
_LmHead? _lm;
|
||
|
||
/// 주 그래프의 W (전치본을 만들 원본): raw 위치 · 크기 · 형
|
||
int? _wRaw, _wRows, _wCols, _wType;
|
||
|
||
_Rewriter(this.src, this.data, this.location);
|
||
|
||
/// ModelProto: graph (7) 만 다시 쓰고 나머지는 그대로
|
||
Uint8List model() {
|
||
_analyze();
|
||
final out = BytesBuilder(copy: false);
|
||
_fields(0, src.lengthSync(), (no, wt, keyStart, valStart, valLen, end) {
|
||
if (no == 7 && wt == 2) {
|
||
_bytes(out, 7, _graph(valStart, valLen, main: true));
|
||
} else {
|
||
out.add(_read(keyStart, end - keyStart));
|
||
}
|
||
});
|
||
return out.takeBytes();
|
||
}
|
||
|
||
// ───────── 분석: 정수 lm_head 로 바꿀 곳 찾기 ─────────
|
||
|
||
void _analyze() {
|
||
final mainInits = <String>{};
|
||
final subgraphs = <List<_Node>>[];
|
||
void graph(int start, int len, bool main) {
|
||
final nodes = <_Node>[];
|
||
_fields(start, len, (no, wt, keyStart, valStart, valLen, end) {
|
||
if (no == 1 && wt == 2) {
|
||
nodes.add(_parseNode(valStart, valLen, (gs, gl) => graph(gs, gl, false)));
|
||
} else if (main && no == 5 && wt == 2) {
|
||
_fields(valStart, valLen, (n, w, ks, vs, vl, e) {
|
||
if (n == 8 && w == 2) mainInits.add(utf8.decode(_read(vs, vl)));
|
||
});
|
||
}
|
||
});
|
||
if (!main) subgraphs.add(nodes);
|
||
}
|
||
|
||
_fields(0, src.lengthSync(), (no, wt, keyStart, valStart, valLen, end) {
|
||
if (no == 7 && wt == 2) graph(valStart, valLen, true);
|
||
});
|
||
|
||
_LmHead? found;
|
||
for (final nodes in subgraphs) {
|
||
for (final tr in nodes) {
|
||
if (tr.op != 'Transpose' || tr.inputs.length != 1 || !mainInits.contains(tr.inputs[0])) continue;
|
||
if (tr.perm.isNotEmpty && (tr.perm.length != 2 || tr.perm[0] != 1 || tr.perm[1] != 0)) continue;
|
||
final dq = nodes.where((n) => n.op == 'DequantizeLinear' && n.inputs.length == 3 && n.inputs[0] == tr.outputs[0]);
|
||
if (dq.length != 1) continue;
|
||
final mm = nodes.where((n) => n.op == 'MatMul' && n.inputs.length == 2 && n.inputs[1] == dq.first.outputs[0]);
|
||
if (mm.length != 1) continue;
|
||
final lm = _LmHead(tr.inputs[0], tr.outputs[0], dq.first.outputs[0], dq.first.inputs[1], dq.first.inputs[2],
|
||
mm.first.inputs[0], mm.first.outputs[0]);
|
||
// 여러 갈래에 같은 모양으로 있어야 한다 (NLLB merged 디코더: If 의 두 갈래)
|
||
if (found != null && found.key != lm.key) return;
|
||
found = lm;
|
||
}
|
||
}
|
||
_lm = found;
|
||
}
|
||
|
||
_Node _parseNode(int start, int len, void Function(int, int) subgraph) {
|
||
final n = _Node();
|
||
_fields(start, len, (no, wt, keyStart, valStart, valLen, end) {
|
||
if (wt != 2) return;
|
||
switch (no) {
|
||
case 1:
|
||
n.inputs.add(utf8.decode(_read(valStart, valLen)));
|
||
case 2:
|
||
n.outputs.add(utf8.decode(_read(valStart, valLen)));
|
||
case 4:
|
||
n.op = utf8.decode(_read(valStart, valLen));
|
||
case 5: // AttributeProto
|
||
var name = '';
|
||
final ints = <int>[];
|
||
_fields(valStart, valLen, (an, aw, aks, avs, avl, ae) {
|
||
if (an == 1 && aw == 2) name = utf8.decode(_read(avs, avl));
|
||
if (an == 6 && aw == 2) subgraph(avs, avl); // g
|
||
if (an == 11 && aw == 2) subgraph(avs, avl); // graphs
|
||
if (an == 8 && aw == 0) ints.add(_readVarint(avs).$1);
|
||
if (an == 8 && aw == 2) {
|
||
for (var q = avs; q < avs + avl;) {
|
||
final (v, k) = _readVarint(q);
|
||
ints.add(v);
|
||
q += k;
|
||
}
|
||
}
|
||
});
|
||
if (name == 'perm') n.perm = ints;
|
||
}
|
||
});
|
||
return n;
|
||
}
|
||
|
||
// ───────── 다시 쓰기 ─────────
|
||
|
||
/// GraphProto: initializer (5) 를 외부 데이터로, 노드 (1) 안의 하위 그래프를 다시 쓰고,
|
||
/// 정수 lm_head 로 바꾸는 갈래에서는 세 노드를 빼고 새 노드를 붙인다
|
||
Uint8List _graph(int start, int len, {required bool main}) {
|
||
final out = BytesBuilder(copy: false);
|
||
final lm = _lm;
|
||
var replaced = false;
|
||
_fields(start, len, (no, wt, keyStart, valStart, valLen, end) {
|
||
if (no == 5 && wt == 2) {
|
||
_bytes(out, 5, _tensor(valStart, valLen, main: main));
|
||
} else if (no == 1 && wt == 2) {
|
||
final n = _parseNode(valStart, valLen, (_, _) {});
|
||
if (lm != null &&
|
||
((n.op == 'Transpose' && n.outputs.contains(lm.transposed)) ||
|
||
(n.op == 'DequantizeLinear' && n.outputs.contains(lm.dequantized)) ||
|
||
(n.op == 'MatMul' && n.outputs.contains(lm.output) && n.inputs.contains(lm.dequantized)))) {
|
||
replaced = true;
|
||
return; // 뺀다
|
||
}
|
||
_bytes(out, 1, _node(valStart, valLen));
|
||
} else {
|
||
out.add(_read(keyStart, end - keyStart));
|
||
}
|
||
});
|
||
if (replaced && lm != null) {
|
||
// 정수 lm_head: out = Cast(MatMulInteger(DQL(h), Wᵀ)) × (h 배율 × W 배율)
|
||
for (final nb in [
|
||
_newNode('DynamicQuantizeLinear', [lm.input], ['jj_lm_hq', 'jj_lm_hs', 'jj_lm_hz'], 'jj_lm_dq'),
|
||
_newNode('MatMulInteger', ['jj_lm_hq', lm.transposed, 'jj_lm_hz', lm.zeroPoint], ['jj_lm_i32'], 'jj_lm_mmi'),
|
||
_newNode('Cast', ['jj_lm_i32'], ['jj_lm_f'], 'jj_lm_cast', castTo: 1), // FLOAT
|
||
_newNode('Mul', ['jj_lm_hs', lm.scale], ['jj_lm_s'], 'jj_lm_scale'),
|
||
_newNode('Mul', ['jj_lm_f', 'jj_lm_s'], [lm.output], 'jj_lm_out'),
|
||
]) {
|
||
_bytes(out, 1, nb);
|
||
}
|
||
}
|
||
// 주 그래프 끝에 전치한 가중치 Wᵀ 를 더한다 (W 는 단어 찾기 (Gather) 에 계속 쓰인다)
|
||
if (main && lm != null && _wRaw != null) {
|
||
_bytes(out, 5, _transposedInit(lm.transposed));
|
||
}
|
||
return out.takeBytes();
|
||
}
|
||
|
||
/// NodeProto: 속성 (5) 안의 하위 그래프만 다시 쓴다
|
||
Uint8List _node(int start, int len) {
|
||
final out = BytesBuilder(copy: false);
|
||
_fields(start, len, (no, wt, keyStart, valStart, valLen, end) {
|
||
if (no == 5 && wt == 2) {
|
||
_bytes(out, 5, _attribute(valStart, valLen));
|
||
} else {
|
||
out.add(_read(keyStart, end - keyStart));
|
||
}
|
||
});
|
||
return out.takeBytes();
|
||
}
|
||
|
||
/// AttributeProto: g (6) · graphs (11)
|
||
Uint8List _attribute(int start, int len) {
|
||
final out = BytesBuilder(copy: false);
|
||
_fields(start, len, (no, wt, keyStart, valStart, valLen, end) {
|
||
if ((no == 6 || no == 11) && wt == 2) {
|
||
_bytes(out, no, _graph(valStart, valLen, main: false));
|
||
} else {
|
||
out.add(_read(keyStart, end - keyStart));
|
||
}
|
||
});
|
||
return out.takeBytes();
|
||
}
|
||
|
||
/// TensorProto: 큰 raw_data (9) 를 .data 로 옮기고 external_data (13) · data_location (14) 를 붙인다
|
||
Uint8List _tensor(int start, int len, {required bool main}) {
|
||
final out = BytesBuilder(copy: false);
|
||
int? rawStart, rawLen, type;
|
||
var name = '';
|
||
final dims = <int>[];
|
||
var external = false;
|
||
_fields(start, len, (no, wt, keyStart, valStart, valLen, end) {
|
||
if (no == 14 && wt == 0) external = true; // 이미 외부 데이터
|
||
if (no == 8 && wt == 2) name = utf8.decode(_read(valStart, valLen));
|
||
if (no == 2 && wt == 0) type = _readVarint(valStart).$1;
|
||
if (no == 1 && wt == 0) dims.add(_readVarint(valStart).$1);
|
||
if (no == 1 && wt == 2) {
|
||
for (var q = valStart; q < valStart + valLen;) {
|
||
final (v, k) = _readVarint(q);
|
||
dims.add(v);
|
||
q += k;
|
||
}
|
||
}
|
||
if (no == 9 && wt == 2 && valLen >= OnnxExternal.threshold) {
|
||
rawStart = valStart;
|
||
rawLen = valLen;
|
||
} else {
|
||
out.add(_read(keyStart, end - keyStart));
|
||
}
|
||
});
|
||
if (rawStart == null || external) {
|
||
// 옮길 것이 없으면 원래 그대로
|
||
if (rawStart != null) return _read(start, len);
|
||
return out.takeBytes();
|
||
}
|
||
// 전치본을 만들 W (int8 · uint8 2차원)
|
||
if (main && name == _lm?.weight && dims.length == 2 && (type == 2 || type == 3) && dims[0] * dims[1] == rawLen) {
|
||
_wRaw = rawStart;
|
||
_wRows = dims[0];
|
||
_wCols = dims[1];
|
||
_wType = type;
|
||
}
|
||
final offset = _alignData();
|
||
const chunk = 4 << 20;
|
||
for (var done = 0; done < rawLen!; done += chunk) {
|
||
final n = rawLen! - done < chunk ? rawLen! - done : chunk;
|
||
data.writeFromSync(_read(rawStart! + done, n));
|
||
}
|
||
_externalFields(out, offset, rawLen!);
|
||
return out.takeBytes();
|
||
}
|
||
|
||
/// W [rows × cols] 를 전치해 .data 에 쓰고, Wᵀ initializer 를 만든다.
|
||
/// 256 열씩 나눠 (64MB) W 를 여러 번 읽는다 - 큰 메모리를 한 번에 잡지 않도록.
|
||
Uint8List _transposedInit(String name) {
|
||
final rows = _wRows!, cols = _wCols!, raw = _wRaw!;
|
||
final offset = _alignData();
|
||
const block = 256, readRows = 4096;
|
||
for (var c0 = 0; c0 < cols; c0 += block) {
|
||
final bc = cols - c0 < block ? cols - c0 : block;
|
||
final outBuf = Uint8List(bc * rows);
|
||
for (var r0 = 0; r0 < rows; r0 += readRows) {
|
||
final br = rows - r0 < readRows ? rows - r0 : readRows;
|
||
final inBuf = _read(raw + r0 * cols, br * cols);
|
||
for (var r = 0; r < br; r++) {
|
||
final ib = r * cols + c0;
|
||
final ob = r0 + r;
|
||
for (var c = 0; c < bc; c++) {
|
||
outBuf[c * rows + ob] = inBuf[ib + c];
|
||
}
|
||
}
|
||
}
|
||
data.writeFromSync(outBuf);
|
||
}
|
||
final out = BytesBuilder(copy: false);
|
||
_key(out, 1, 0);
|
||
_varint(out, cols);
|
||
_key(out, 1, 0);
|
||
_varint(out, rows);
|
||
_key(out, 2, 0);
|
||
_varint(out, _wType!);
|
||
_str(out, 8, name);
|
||
_externalFields(out, offset, rows * cols);
|
||
return out.takeBytes();
|
||
}
|
||
|
||
/// .data 끝을 정렬 위치까지 0 으로 채우고 그 위치를 돌려준다
|
||
int _alignData() {
|
||
var offset = data.lengthSync();
|
||
final pad = (OnnxExternal.align - offset % OnnxExternal.align) % OnnxExternal.align;
|
||
data.setPositionSync(offset);
|
||
if (pad > 0) {
|
||
data.writeFromSync(Uint8List(pad));
|
||
offset += pad;
|
||
}
|
||
return offset;
|
||
}
|
||
|
||
void _externalFields(BytesBuilder out, int offset, int length) {
|
||
for (final (k, v) in [('location', location), ('offset', '$offset'), ('length', '$length')]) {
|
||
final entry = BytesBuilder(copy: false);
|
||
_str(entry, 1, k);
|
||
_str(entry, 2, v);
|
||
_bytes(out, 13, entry.takeBytes());
|
||
}
|
||
_key(out, 14, 0);
|
||
_varint(out, 1); // EXTERNAL
|
||
}
|
||
|
||
/// 새 NodeProto (기본 도메인). [castTo]: Cast 의 to 속성
|
||
static Uint8List _newNode(String op, List<String> inputs, List<String> outputs, String name, {int? castTo}) {
|
||
final out = BytesBuilder(copy: false);
|
||
for (final i in inputs) {
|
||
_str(out, 1, i);
|
||
}
|
||
for (final o in outputs) {
|
||
_str(out, 2, o);
|
||
}
|
||
_str(out, 3, name);
|
||
_str(out, 4, op);
|
||
if (castTo != null) {
|
||
final a = BytesBuilder(copy: false);
|
||
_str(a, 1, 'to');
|
||
_key(a, 3, 0);
|
||
_varint(a, castTo);
|
||
_key(a, 20, 0);
|
||
_varint(a, 2); // AttributeType.INT
|
||
_bytes(out, 5, a.takeBytes());
|
||
}
|
||
return out.takeBytes();
|
||
}
|
||
|
||
/// [start] 부터 [len] 바이트 안의 protobuf 필드를 차례로
|
||
void _fields(int start, int len,
|
||
void Function(int no, int wt, int keyStart, int valStart, int valLen, int end) f) {
|
||
var pos = start;
|
||
final end = start + len;
|
||
while (pos < end) {
|
||
final keyStart = pos;
|
||
final (key, kn) = _readVarint(pos);
|
||
pos += kn;
|
||
final no = key >> 3, wt = key & 7;
|
||
int valStart = pos, valLen;
|
||
switch (wt) {
|
||
case 0:
|
||
final (_, n) = _readVarint(pos);
|
||
valLen = n;
|
||
case 1:
|
||
valLen = 8;
|
||
case 5:
|
||
valLen = 4;
|
||
case 2:
|
||
final (l, n) = _readVarint(pos);
|
||
valStart = pos + n;
|
||
valLen = l;
|
||
default:
|
||
throw FormatException(trf('지원하지 않는 protobuf 형식 (wire type {0}, 위치 {1})', [wt, keyStart]));
|
||
}
|
||
final fieldEnd = valStart + valLen;
|
||
if (fieldEnd > end) throw FormatException(trf('모델 파일이 잘렸습니다 (위치 {0})', [keyStart]));
|
||
f(no, wt, keyStart, valStart, valLen, fieldEnd);
|
||
pos = fieldEnd;
|
||
}
|
||
}
|
||
|
||
(int, int) _readVarint(int pos) {
|
||
src.setPositionSync(pos);
|
||
final n = src.readIntoSync(_buf);
|
||
var v = 0, shift = 0;
|
||
for (var i = 0; i < n; i++) {
|
||
final b = _buf[i];
|
||
v |= (b & 0x7f) << shift;
|
||
if (b < 0x80) return (v, i + 1);
|
||
shift += 7;
|
||
}
|
||
throw FormatException(trf('잘못된 varint (위치 {0})', [pos]));
|
||
}
|
||
|
||
Uint8List _read(int pos, int len) {
|
||
src.setPositionSync(pos);
|
||
final b = src.readSync(len);
|
||
if (b.length != len) throw FormatException(trf('모델 파일이 잘렸습니다 (위치 {0})', [pos]));
|
||
return b;
|
||
}
|
||
|
||
static void _key(BytesBuilder out, int no, int wt) => _varint(out, (no << 3) | wt);
|
||
|
||
static void _varint(BytesBuilder out, int v) {
|
||
while (v >= 0x80) {
|
||
out.addByte((v & 0x7f) | 0x80);
|
||
v >>= 7;
|
||
}
|
||
out.addByte(v);
|
||
}
|
||
|
||
static void _bytes(BytesBuilder out, int no, Uint8List b) {
|
||
_key(out, no, 2);
|
||
_varint(out, b.length);
|
||
out.add(b);
|
||
}
|
||
|
||
static void _str(BytesBuilder out, int no, String s) => _bytes(out, no, utf8.encode(s));
|
||
}
|