-
-
Notifications
You must be signed in to change notification settings - Fork 410
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Update performance and reduce memory usage for big Labels layer in di…
…rect color mode (#6439) Closes #6518 Closes #6084 # Description In this PR, similarly to #6411, instead of using `float32` to pass data to the GPU there we introduce heuristics for choosing smaller data types, while keeping high performance. Instead of complex calculation of color in the shader, a precomputed texture array is used. To avoid repetitive texture calculation, the textures are cached in the `Colormap` objects. For data of type uint8/int8/uint16/int16 we do not perform any transform of data. We send them to the GPU as it is. This allows to reduce computational time. Based on experiments, the rendering performance is a little worse for uint16/int16 than for uint8/int8. But it may depend on the GPU. Also, using uint16/int16 means usage more GPU memory than for 8 bits type. Still less than current main. For datatypes using at least 32 bits, we add a preprocessing step where we identify a set of labels that are mapped to the same color and map all of them to the same value. This often saves enough space to fall back to uint8/uint16. It allows using a smaller additional array, and use less GPU memory. If there are more than `2**16` distinct colors, then float32 is used, though performance will be reduced. We support only up to `2**23` distinct colors for now. For reduced memory usage, part of the functions used for data preprocessing are compiled using numba. We provide a version of the function that does not require `numba` but it limits the number of distinct colors to `2**16` and involves additional array creation (more memory usage). --------- Co-authored-by: Juan Nunez-Iglesias <jni@fastmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lorenzo Gaifas <brisvag@gmail.com> Co-authored-by: Andy Sweet <andrew.d.sweet@gmail.com>
- Loading branch information
5 people
committed
Dec 15, 2023
1 parent
a45683f
commit 6524ee4
Showing
13 changed files
with
1,122 additions
and
794 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,121 +1,24 @@ | ||
| from itertools import product | ||
| from unittest.mock import patch | ||
|
|
||
| import numpy as np | ||
| import pytest | ||
|
|
||
| from napari._vispy.layers.labels import ( | ||
| MAX_LOAD_FACTOR, | ||
| PRIME_NUM_TABLE, | ||
| build_textures_from_dict, | ||
| hash2d_get, | ||
| idx_to_2d, | ||
| ) | ||
|
|
||
|
|
||
| @pytest.fixture(scope='module', autouse=True) | ||
| def mock_max_texture_size(): | ||
| """When running tests in this file, pretend max texture size is 2^16.""" | ||
| with patch('napari._vispy.layers.labels.MAX_TEXTURE_SIZE', 2**16): | ||
| yield | ||
|
|
||
|
|
||
| def test_idx_to_2d(): | ||
| assert idx_to_2d(0, (100, 100)) == (0, 0) | ||
| assert idx_to_2d(1, (100, 100)) == (0, 1) | ||
| assert idx_to_2d(101, (100, 100)) == (1, 1) | ||
| assert idx_to_2d(521, (100, 100)) == (5, 21) | ||
| assert idx_to_2d(100 * 100 + 521, (100, 100)) == (5, 21) | ||
|
|
||
|
|
||
| def test_build_textures_from_dict(): | ||
| keys, values, collision = build_textures_from_dict( | ||
| {1: (1, 1, 1, 1), 2: (2, 2, 2, 2)} | ||
| values = build_textures_from_dict( | ||
| {0: (0, 0, 0, 0), 1: (1, 1, 1, 1), 2: (2, 2, 2, 2)}, | ||
| max_size=10, | ||
| ) | ||
| assert not collision | ||
| assert keys.shape == (37, 37) | ||
| assert values.shape == (37, 37, 4) | ||
| assert keys[0, 1] == 1 | ||
| assert keys[0, 2] == 2 | ||
| assert np.array_equiv(values[0, 1], (1, 1, 1, 1)) | ||
| assert np.array_equiv(values[0, 2], (2, 2, 2, 2)) | ||
| assert values.shape == (3, 1, 4) | ||
| assert np.array_equiv(values[1], (1, 1, 1, 1)) | ||
| assert np.array_equiv(values[2], (2, 2, 2, 2)) | ||
|
|
||
|
|
||
| def test_build_textures_from_dict_too_many_labels(monkeypatch): | ||
| with pytest.raises(MemoryError): | ||
| build_textures_from_dict( | ||
| {i: (i, i, i, i) for i in range(1001)}, shape=(10, 10) | ||
| ) | ||
| monkeypatch.setattr( | ||
| "napari._vispy.layers.labels.PRIME_NUM_TABLE", [[61], [127]] | ||
| ) | ||
| with pytest.raises(MemoryError): | ||
| def test_build_textures_from_dict_exc(): | ||
| with pytest.raises(ValueError, match="Cannot create a 2D texture"): | ||
| build_textures_from_dict( | ||
| {i: (i, i, i, i) for i in range((251**2) // 2)}, | ||
| {0: (0, 0, 0, 0), 1: (1, 1, 1, 1), 2: (2, 2, 2, 2)}, | ||
| max_size=1, | ||
| ) | ||
|
|
||
|
|
||
| def test_size_of_texture_square(): | ||
| count = int(127 * 127 * MAX_LOAD_FACTOR) - 1 | ||
| keys, values, *_ = build_textures_from_dict( | ||
| {i: (i, i, i, i) for i in range(count)} | ||
| ) | ||
| assert keys.shape == (127, 127) | ||
| assert values.shape == (127, 127, 4) | ||
|
|
||
|
|
||
| def test_size_of_texture_rectangle(): | ||
| count = int(128 * 128 * MAX_LOAD_FACTOR) + 5 | ||
| keys, values, *_ = build_textures_from_dict( | ||
| {i: (i, i, i, i) for i in range(count)} | ||
| ) | ||
| assert keys.shape == (251, 127) | ||
| assert values.shape == (251, 127, 4) | ||
|
|
||
|
|
||
| def test_build_textures_from_dict_collision(): | ||
| keys, values, collision = build_textures_from_dict( | ||
| {1: (1, 1, 1, 1), 26: (2, 2, 2, 2), 27: (3, 3, 3, 3)}, shape=(5, 5) | ||
| ) | ||
| assert collision | ||
| assert keys.shape == (5, 5) | ||
| assert keys[0, 1] == 1 | ||
| assert keys[0, 2] == 26 | ||
| assert keys[0, 3] == 27 | ||
| assert np.array_equiv(values[0, 1], (1, 1, 1, 1)) | ||
| assert np.array_equiv(values[0, 2], (2, 2, 2, 2)) | ||
| assert np.array_equiv(values[0, 3], (3, 3, 3, 3)) | ||
|
|
||
| assert hash2d_get(1, keys) == (0, 1) | ||
| assert hash2d_get(26, keys) == (0, 2) | ||
| assert hash2d_get(27, keys) == (0, 3) | ||
|
|
||
|
|
||
| def test_collide_keys(): | ||
| base_keys = [x * y for x, y in product(PRIME_NUM_TABLE[0], repeat=2)] | ||
| colors = {0: (0, 0, 0, 0), 1: (1, 1, 1, 1)} | ||
| colors.update({i + 10: (1, 0, 0, 1) for i in base_keys}) | ||
| colors.update({2 * i + 10: (0, 1, 0, 1) for i in base_keys}) | ||
| keys, values, collision = build_textures_from_dict(colors) | ||
| assert not collision | ||
| assert keys.shape == (37, 61) | ||
| assert values.shape == (37, 61, 4) | ||
|
|
||
|
|
||
| def test_collide_keys2(): | ||
| base_keys = [x * y for x, y in product(PRIME_NUM_TABLE[0], repeat=2)] + [ | ||
| x * y for x, y in product(PRIME_NUM_TABLE[0], PRIME_NUM_TABLE[1]) | ||
| ] | ||
| colors = {0: (0, 0, 0, 0), 1: (1, 1, 1, 1)} | ||
| colors.update({i + 10: (1, 0, 0, 1) for i in base_keys}) | ||
| colors.update({2 * i + 10: (0, 1, 0, 1) for i in base_keys}) | ||
|
|
||
| # enforce collision for collision table of size 31 | ||
| colors.update({31 * i + 10: (0, 0, 1, 1) for i in base_keys}) | ||
| # enforce collision for collision table of size 29 | ||
| colors.update({29 * i + 10: (0, 0, 1, 1) for i in base_keys}) | ||
|
|
||
| keys, values, collision = build_textures_from_dict(colors) | ||
| assert collision | ||
| assert keys.shape == (37, 37) | ||
| assert values.shape == (37, 37, 4) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.