diff --git a/README.md b/README.md index 7321780..0d44a35 100644 --- a/README.md +++ b/README.md @@ -98,9 +98,14 @@ compiled.dump('compiled.pkl') restored_compiled = TLCompiledPredictor.load('compiled.pkl') ``` +Threshold quantization can be enabled by passing `params={'quantize': 1}` +to `compile()`; prediction keeps the original float32 split boundaries. + `TLCompiledPredictor.predict()` also accepts SciPy sparse matrices (CSR, CSC and COO) without converting them to dense arrays. Unstored entries and stored NaNs are treated as missing values; explicitly stored zeros remain zero. +Duplicate coordinates are summed in the input dtype before conversion to +float32. Prediction leaves the input matrix unchanged. Compiled predictor dumps contain metadata and the library path. Deploy the `.so`/`.dll` alongside the dump. If its path changes, load it with diff --git a/py_boost/utils/tl_wrapper.py b/py_boost/utils/tl_wrapper.py index b34e906..492d2e7 100644 --- a/py_boost/utils/tl_wrapper.py +++ b/py_boost/utils/tl_wrapper.py @@ -99,7 +99,13 @@ def _prepare_tl2cgen_input(X): # The new sparse object may share buffers with X; replace rather than mutate # those buffers. TL2cgen treats unstored entries as missing, but does not # recognize NaNs stored explicitly in a CSR buffer. Drop only NaNs, not zeros. - matrix = csr_matrix(X, dtype=np.float32, copy=False) + matrix = csr_matrix(X, copy=False) + if not matrix.has_canonical_format: + matrix = matrix.copy() + matrix.sum_duplicates() + # Sum in the input dtype before casting, as COO conversion already does. + # In particular, NaN plus another stored value must remain missing. + matrix = matrix.astype(np.float32, copy=False) missing = np.isnan(matrix.data) if missing.any(): missing_rows = np.searchsorted(matrix.indptr, np.flatnonzero(missing), side='right') - 1 @@ -116,6 +122,36 @@ def _prepare_tl2cgen_input(X): return matrix +def _prepare_tl2cgen_model(model): + """Preserve float32 split boundaries in TL2cgen's eight-digit quantizer.""" + if model.input_type != 'float32': + return model + prepared = None + for tree_id in range(model.num_tree): + tree = model.get_tree_accessor(tree_id) + thresholds = tree.get_field('threshold') + comparison = tree.get_field('cmp') + # Treelite's Operator enum uses kLE=3 and kLT=2. + candidates = np.flatnonzero((comparison == 3) & np.isfinite(thresholds)) + affected = [i for i in candidates + if np.float32(format(thresholds[i], '.8g')) != thresholds[i]] + if not affected: + continue + if prepared is None: + prepared = treelite.Model.deserialize_bytes(model.serialize_bytes()) + thresholds, comparison = thresholds.copy(), comparison.copy() + # Adjacent float32 values have no representable value between them: + # x <= t is equivalent to x < nextafter(t, +inf). If t fails the + # eight-digit round trip, its successor round-trips exactly. Change + # only the compilation copy, including when loading old checkpoints. + thresholds[affected] = np.nextafter(thresholds[affected], np.float32(np.inf)) + comparison[affected] = 2 + tree = prepared.get_tree_accessor(tree_id) + tree.set_field('threshold', thresholds) + tree.set_field('cmp', comparison) + return model if prepared is None else prepared + + class TLCompiledPredictor: """ Compiled TL2cgen model for CPU inference @@ -269,9 +305,12 @@ def compile( raise ValueError("TL2cgen only supports compiler='ast_native'") params = {} if params is None else params params = {**{'parallel_comp': os.cpu_count() or 1}, **params} + model = self.tl_model + if params.get('quantize', 0) > 0: + model = _prepare_tl2cgen_model(model) tl2cgen.export_lib( - self.tl_model, toolchain=toolchain, libpath=libpath, params=params, + model, toolchain=toolchain, libpath=libpath, params=params, verbose=verbose, nthread=nthread, options=options, ) diff --git a/tests/test_treelite.py b/tests/test_treelite.py index 58ba8ed..d7cff12 100644 --- a/tests/test_treelite.py +++ b/tests/test_treelite.py @@ -147,6 +147,90 @@ def no_dense(*args, **kwargs): np.testing.assert_array_equal(actual_buffer, original_buffer) +@pytest.mark.parametrize('format', ['csr', 'csc', 'coo']) +@pytest.mark.parametrize('dtype', [np.float32, np.float64]) +@pytest.mark.parametrize('index_dtype', [np.int32, np.int64]) +def test_compiled_sparse_duplicate_coordinates(ensemble, sparse_compiled, format, dtype, + index_dtype, monkeypatch): + # Include unsorted duplicates, a sum crossing a split, cancellation to an + # explicitly stored zero, and NaN in either position among duplicates. + large = 1e8 if dtype == np.float64 else 1e4 + X = sparse.csr_matrix(( + np.array([.4, -1, .4, np.nan, -1, 1, 1, -1, np.nan, + .25, .25, -.25, -.25, np.nan, 0, -1, .2, 0, .2, + large, .75, -large, -1], dtype=dtype), + [0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 0, 1], + [0, 3, 6, 9, 13, 16, 19, 23], + ), shape=(7, 2)).asformat(format) + if format == 'coo': + X.row = X.row.astype(index_dtype) + X.col = X.col.astype(index_dtype) + fields = ('row', 'col') + else: + X.indices = X.indices.astype(index_dtype) + X.indptr = X.indptr.astype(index_dtype) + assert not X.has_canonical_format + fields = ('indices', 'indptr') + original = copy.deepcopy(X) + + def no_dense(*args, **kwargs): + pytest.fail('Sparse prediction must not densify its input') + + for cls in (sparse.csr_matrix, sparse.csc_matrix, sparse.coo_matrix): + monkeypatch.setattr(cls, 'toarray', no_dense) + + expected = np.array([ + [-4, 4, 7], [2, -4, 5.5], [2, -4, 5.5], + [2, -4, 5.5], [4, 2, 4], [4, 2, 4], [-4, 4, 7], + ], dtype=np.float32)[:, :len(ensemble.base_score)] + np.testing.assert_allclose(sparse_compiled.predict(X), expected) + np.testing.assert_array_equal(X.data, original.data) + for field in fields: + actual_buffer, original_buffer = getattr(X, field), getattr(original, field) + assert actual_buffer.dtype == original_buffer.dtype + np.testing.assert_array_equal(actual_buffer, original_buffer) + if format != 'coo': + assert not X.has_canonical_format + + +@pytest.mark.parametrize('threshold', [-0.11849978566169739, 0.11849978566169739]) +@pytest.mark.parametrize('nan_left', [False, True]) +@pytest.mark.parametrize('parallel_comp', [0, 2]) +def test_quantized_threshold_equality(ensemble, tmp_path, toolchain, threshold, + nan_left, parallel_comp): + threshold = np.float32(threshold) + # TL2cgen 1.0 rounds these float32 thresholds to a different value when + # emitting its quantizer array with only eight significant decimal digits. + assert np.float32(format(threshold, '.8g')) != threshold + model = CPUEnsemble(len(ensemble.base_score)) + model.models[0].test_format = np.array([ + -1 if nan_left else 1, threshold, -1, -2, + ], dtype=np.float32) + predictor = TLPredictor(model) + # Exercise old checkpoints too: the fix must run at compilation time. + predictor.dump(tmp_path / 'gtil') + predictor = TLPredictor.load(tmp_path / 'gtil') + before = predictor.tl_model.dump_as_json() + name, suffix = toolchain + params = {'quantize': 1, 'parallel_comp': parallel_comp} + compiled = predictor.compile(name, tmp_path / ('quantized' + suffix), + params=params, nthread=2) + assert params == {'quantize': 1, 'parallel_comp': parallel_comp} + assert predictor.tl_model.dump_as_json() == before + + X = np.zeros((6, 2), dtype=np.float32) + X[:, 0] = [np.nextafter(threshold, np.float32(-np.inf)), threshold, + np.nextafter(threshold, np.float32(np.inf)), np.nan, -np.inf, np.inf] + left, right = [2, -4, 5.5], [4, 2, 4] + expected = np.array([left, left, right, left if nan_left else right, left, right], + dtype=np.float32)[:, :len(ensemble.base_score)] + np.testing.assert_allclose(predictor.predict(X), expected) + for data in (X, X.astype(np.float64), sparse.csr_matrix(X)): + actual = compiled.predict(data) + assert actual.dtype == np.float32 + np.testing.assert_allclose(actual, expected) + + def test_shallow_copy_keeps_native_predictors(ensemble, compiled, examples): X, expected = examples expected = expected[:, :len(ensemble.base_score)]