enczoo 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- enczoo-0.1.0/LICENSE +21 -0
- enczoo-0.1.0/PKG-INFO +74 -0
- enczoo-0.1.0/README.md +37 -0
- enczoo-0.1.0/pyproject.toml +32 -0
- enczoo-0.1.0/src/enczoo/__init__.py +6 -0
- enczoo-0.1.0/src/enczoo/base.py +172 -0
- enczoo-0.1.0/src/enczoo/classic/__init__.py +0 -0
- enczoo-0.1.0/src/enczoo/classic/pixels.py +89 -0
- enczoo-0.1.0/src/enczoo/neural_networks/__init__.py +0 -0
- enczoo-0.1.0/src/enczoo/neural_networks/base.py +184 -0
- enczoo-0.1.0/src/enczoo/neural_networks/pretrained.py +163 -0
- enczoo-0.1.0/src/enczoo/neural_networks/tests/__init__.py +0 -0
- enczoo-0.1.0/src/enczoo/neural_networks/tests/test_pretrained_nn.py +54 -0
- enczoo-0.1.0/src/enczoo/py.typed +0 -0
- enczoo-0.1.0/src/enczoo/transforms/__init__.py +0 -0
- enczoo-0.1.0/src/enczoo/transforms/random_projection.py +73 -0
- enczoo-0.1.0/src/enczoo/utils.py +60 -0
enczoo-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Michael J. Lee
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
enczoo-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: enczoo
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Map images (as `PIL.Images`) to intermediate representations (as `np.ndarray`) from off-the-shelf vision models.
|
|
5
|
+
Author: himjl
|
|
6
|
+
Author-email: himjl <mil@mit.edu>
|
|
7
|
+
License: MIT License
|
|
8
|
+
|
|
9
|
+
Copyright (c) 2026 Michael J. Lee
|
|
10
|
+
|
|
11
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
12
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
13
|
+
in the Software without restriction, including without limitation the rights
|
|
14
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
15
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
16
|
+
furnished to do so, subject to the following conditions:
|
|
17
|
+
|
|
18
|
+
The above copyright notice and this permission notice shall be included in all
|
|
19
|
+
copies or substantial portions of the Software.
|
|
20
|
+
|
|
21
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
22
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
23
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
24
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
25
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
26
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
27
|
+
SOFTWARE.
|
|
28
|
+
Requires-Dist: h5py>=3.15.1
|
|
29
|
+
Requires-Dist: pillow>=12.1.0
|
|
30
|
+
Requires-Dist: pydantic>=2.12.5
|
|
31
|
+
Requires-Dist: pytest>=9.0.2
|
|
32
|
+
Requires-Dist: torch>=2.9.1
|
|
33
|
+
Requires-Dist: torchvision>=0.24.1
|
|
34
|
+
Requires-Dist: tqdm>=4.67.1
|
|
35
|
+
Requires-Python: >=3.12
|
|
36
|
+
Description-Content-Type: text/markdown
|
|
37
|
+
|
|
38
|
+
# `enczoo`: a zoo of encoding models for images
|
|
39
|
+
|
|
40
|
+
[](https://github.com/himjl/enczoo/actions/workflows/ci.yml)
|
|
41
|
+
|
|
42
|
+
`enczoo` is a Python library with a single goal: to map images (as `PIL.Images`) to intermediate representations (as `np.ndarray`) from off-the-shelf vision models, such as AlexNet and ResNet50.
|
|
43
|
+
|
|
44
|
+
This library is meant for those who just need to compute off-the-shelf image features once for their project (and perhaps cache them elsewhere).
|
|
45
|
+
|
|
46
|
+
### Installation
|
|
47
|
+
|
|
48
|
+
`enczoo` requires Python 3.12 or above, and may be installed using [uv](https://docs.astral.sh/uv/) with the following command:
|
|
49
|
+
|
|
50
|
+
>`uv add enczoo`
|
|
51
|
+
|
|
52
|
+
### Usage
|
|
53
|
+
|
|
54
|
+
```python
|
|
55
|
+
import enczoo
|
|
56
|
+
import PIL.Image
|
|
57
|
+
image = PIL.Image.open('my-image.png')
|
|
58
|
+
|
|
59
|
+
model = enczoo.ResNet50(layer_name='avgpool') # try layer4, layer3, ...
|
|
60
|
+
features = model.compute_features(images=[image]) # np.ndarray
|
|
61
|
+
|
|
62
|
+
# Want another layer? Check out: print(enczoo.ResNet50.layer_names)
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
### Things `enczoo` handles
|
|
66
|
+
`enczoo` aims to "just work" by solving several tiny problems which collectively make computing image features a bit annoying. `enczoo` handles:
|
|
67
|
+
|
|
68
|
+
* performing model-specific image normalization ("_was it -1 to 1, 0 to 1, 0-255...? ImageNet channel normalization...?_"),
|
|
69
|
+
* correctly encoding images ("_my image was in mode L, not RGB!_")
|
|
70
|
+
* turning off any batch normalization ("_was the model in training mode...?_")
|
|
71
|
+
* extracting intermediate layers by name ("_how do I do that forward hook thing again...?_")
|
|
72
|
+
* turning off autograd, and returning tensors as `np.ndarray` (no more `.cpu().numpy()`)
|
|
73
|
+
* image cropping to fit input tensor shape (default: center cropping. no black bars!)
|
|
74
|
+
* and more!
|
enczoo-0.1.0/README.md
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
# `enczoo`: a zoo of encoding models for images
|
|
2
|
+
|
|
3
|
+
[](https://github.com/himjl/enczoo/actions/workflows/ci.yml)
|
|
4
|
+
|
|
5
|
+
`enczoo` is a Python library with a single goal: to map images (as `PIL.Images`) to intermediate representations (as `np.ndarray`) from off-the-shelf vision models, such as AlexNet and ResNet50.
|
|
6
|
+
|
|
7
|
+
This library is meant for those who just need to compute off-the-shelf image features once for their project (and perhaps cache them elsewhere).
|
|
8
|
+
|
|
9
|
+
### Installation
|
|
10
|
+
|
|
11
|
+
`enczoo` requires Python 3.12 or above, and may be installed using [uv](https://docs.astral.sh/uv/) with the following command:
|
|
12
|
+
|
|
13
|
+
>`uv add enczoo`
|
|
14
|
+
|
|
15
|
+
### Usage
|
|
16
|
+
|
|
17
|
+
```python
|
|
18
|
+
import enczoo
|
|
19
|
+
import PIL.Image
|
|
20
|
+
image = PIL.Image.open('my-image.png')
|
|
21
|
+
|
|
22
|
+
model = enczoo.ResNet50(layer_name='avgpool') # try layer4, layer3, ...
|
|
23
|
+
features = model.compute_features(images=[image]) # np.ndarray
|
|
24
|
+
|
|
25
|
+
# Want another layer? Check out: print(enczoo.ResNet50.layer_names)
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
### Things `enczoo` handles
|
|
29
|
+
`enczoo` aims to "just work" by solving several tiny problems which collectively make computing image features a bit annoying. `enczoo` handles:
|
|
30
|
+
|
|
31
|
+
* performing model-specific image normalization ("_was it -1 to 1, 0 to 1, 0-255...? ImageNet channel normalization...?_"),
|
|
32
|
+
* correctly encoding images ("_my image was in mode L, not RGB!_")
|
|
33
|
+
* turning off any batch normalization ("_was the model in training mode...?_")
|
|
34
|
+
* extracting intermediate layers by name ("_how do I do that forward hook thing again...?_")
|
|
35
|
+
* turning off autograd, and returning tensors as `np.ndarray` (no more `.cpu().numpy()`)
|
|
36
|
+
* image cropping to fit input tensor shape (default: center cropping. no black bars!)
|
|
37
|
+
* and more!
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "enczoo"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Map images (as `PIL.Images`) to intermediate representations (as `np.ndarray`) from off-the-shelf vision models."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
authors = [
|
|
7
|
+
{ name = "himjl", email = "mil@mit.edu" }
|
|
8
|
+
]
|
|
9
|
+
license = { file = "LICENSE" }
|
|
10
|
+
requires-python = ">=3.12"
|
|
11
|
+
dependencies = [
|
|
12
|
+
"h5py>=3.15.1",
|
|
13
|
+
"pillow>=12.1.0",
|
|
14
|
+
"pydantic>=2.12.5",
|
|
15
|
+
"pytest>=9.0.2",
|
|
16
|
+
"torch>=2.9.1",
|
|
17
|
+
"torchvision>=0.24.1",
|
|
18
|
+
"tqdm>=4.67.1",
|
|
19
|
+
]
|
|
20
|
+
|
|
21
|
+
[build-system]
|
|
22
|
+
requires = ["uv_build>=0.8.12,<0.9.0"]
|
|
23
|
+
build-backend = "uv_build"
|
|
24
|
+
|
|
25
|
+
[dependency-groups]
|
|
26
|
+
dev = [
|
|
27
|
+
"black>=25.12.0",
|
|
28
|
+
"matplotlib>=3.10.8",
|
|
29
|
+
"pytest>=9.0.2",
|
|
30
|
+
"ruff>=0.14.10",
|
|
31
|
+
"ty>=0.0.9",
|
|
32
|
+
]
|
|
@@ -0,0 +1,172 @@
|
|
|
1
|
+
from abc import ABC, abstractmethod
|
|
2
|
+
from typing import List, Tuple
|
|
3
|
+
|
|
4
|
+
import PIL.Image
|
|
5
|
+
import numpy as np
|
|
6
|
+
import torch
|
|
7
|
+
|
|
8
|
+
import enczoo.utils as utils
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
# %%
|
|
12
|
+
class ImageEncoding(
|
|
13
|
+
torch.nn.Module,
|
|
14
|
+
ABC,
|
|
15
|
+
):
|
|
16
|
+
"""Map PIL images to batched float tensors.
|
|
17
|
+
|
|
18
|
+
This module maps B-length lists of PIL.Image.Image to float tensors shaped
|
|
19
|
+
[B, *]. Parameters do not aggregate gradients by default.
|
|
20
|
+
"""
|
|
21
|
+
|
|
22
|
+
def __init__(
|
|
23
|
+
self,
|
|
24
|
+
trainable: bool = False,
|
|
25
|
+
):
|
|
26
|
+
"""Initialize the encoder.
|
|
27
|
+
|
|
28
|
+
Args:
|
|
29
|
+
trainable: If True, the module is put in train mode.
|
|
30
|
+
"""
|
|
31
|
+
super().__init__()
|
|
32
|
+
self.trainable = trainable
|
|
33
|
+
self._module_hash = None
|
|
34
|
+
self._tensor_bucket = None
|
|
35
|
+
self._output_shape = None
|
|
36
|
+
|
|
37
|
+
# Set the module's mode
|
|
38
|
+
self.train(mode=trainable)
|
|
39
|
+
|
|
40
|
+
@property
|
|
41
|
+
def device(self) -> torch.device:
|
|
42
|
+
"""Infer the device from the first parameter or buffer.
|
|
43
|
+
|
|
44
|
+
Returns:
|
|
45
|
+
The device of the first parameter or buffer, or CPU if none exist.
|
|
46
|
+
"""
|
|
47
|
+
try:
|
|
48
|
+
return next(self.parameters()).device
|
|
49
|
+
except StopIteration:
|
|
50
|
+
try:
|
|
51
|
+
return next(self.buffers()).device
|
|
52
|
+
except StopIteration:
|
|
53
|
+
return torch.device("cpu")
|
|
54
|
+
|
|
55
|
+
@property
|
|
56
|
+
def output_shape(self) -> Tuple[int, ...]:
|
|
57
|
+
"""Return the output feature shape (excluding batch dimension)."""
|
|
58
|
+
if self._output_shape is None:
|
|
59
|
+
test_image = PIL.Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))
|
|
60
|
+
test_result = self.compute_features(images=[test_image], flatten=False)
|
|
61
|
+
if not isinstance(test_result, np.ndarray):
|
|
62
|
+
raise ValueError(
|
|
63
|
+
f"Expected a np.ndarray from self.forward, but got {type(test_result)}"
|
|
64
|
+
)
|
|
65
|
+
if not test_result.shape[0] == 1:
|
|
66
|
+
raise ValueError(
|
|
67
|
+
f"Expected a batch size of 1, but got {test_result.shape}"
|
|
68
|
+
)
|
|
69
|
+
if len(test_result.shape) == 1:
|
|
70
|
+
output_shape = tuple()
|
|
71
|
+
else:
|
|
72
|
+
output_shape = test_result.shape[1:]
|
|
73
|
+
|
|
74
|
+
self._output_shape = tuple(output_shape)
|
|
75
|
+
|
|
76
|
+
return self._output_shape
|
|
77
|
+
|
|
78
|
+
@property
|
|
79
|
+
def module_hash(self) -> str:
|
|
80
|
+
"""Return a stable hash for the module's parameters and structure."""
|
|
81
|
+
if self.trainable:
|
|
82
|
+
raise ValueError("Cannot hash a trainable model.")
|
|
83
|
+
|
|
84
|
+
# Turn off gradients for all parameters
|
|
85
|
+
for param in self.parameters():
|
|
86
|
+
param.requires_grad = False
|
|
87
|
+
|
|
88
|
+
# Hash self if unhashed
|
|
89
|
+
if self._module_hash is None:
|
|
90
|
+
self._module_hash = utils.hash_torch_module(module=self)
|
|
91
|
+
return self._module_hash
|
|
92
|
+
|
|
93
|
+
def compute_features(
|
|
94
|
+
self,
|
|
95
|
+
images: List[PIL.Image.Image],
|
|
96
|
+
flatten: bool = False,
|
|
97
|
+
seed: int | None = None,
|
|
98
|
+
) -> np.ndarray:
|
|
99
|
+
"""Compute features and return them as a NumPy array.
|
|
100
|
+
|
|
101
|
+
This is an alias for __call__ to help IDEs that do not recognize the
|
|
102
|
+
torch.nn.Module __call__ signature.
|
|
103
|
+
|
|
104
|
+
Args:
|
|
105
|
+
images: A B-length list of PIL.Image.Image.
|
|
106
|
+
flatten: If True, flatten the output to [B, d].
|
|
107
|
+
seed: Optional RNG seed for deterministic results.
|
|
108
|
+
|
|
109
|
+
Returns:
|
|
110
|
+
A NumPy array of shape [B, *], or [B, d] if flatten=True.
|
|
111
|
+
|
|
112
|
+
Raises:
|
|
113
|
+
ValueError: If the input images are invalid.
|
|
114
|
+
"""
|
|
115
|
+
|
|
116
|
+
with torch.random.fork_rng():
|
|
117
|
+
if seed is not None:
|
|
118
|
+
torch.manual_seed(seed)
|
|
119
|
+
if torch.cuda.is_available():
|
|
120
|
+
torch.cuda.manual_seed_all(seed)
|
|
121
|
+
with torch.no_grad():
|
|
122
|
+
torch_features: torch.Tensor = self(images=images, flatten=flatten)
|
|
123
|
+
numpy_features = torch_features.detach().cpu().numpy()
|
|
124
|
+
return numpy_features
|
|
125
|
+
|
|
126
|
+
def forward(
|
|
127
|
+
self,
|
|
128
|
+
images: List[PIL.Image.Image],
|
|
129
|
+
flatten: bool = False,
|
|
130
|
+
) -> torch.Tensor:
|
|
131
|
+
"""Compute features for a batch of images.
|
|
132
|
+
|
|
133
|
+
Args:
|
|
134
|
+
images: A B-length list of PIL.Image.Image.
|
|
135
|
+
flatten: If True, flatten the output to [B, d].
|
|
136
|
+
|
|
137
|
+
Returns:
|
|
138
|
+
A torch.Tensor of shape [B, *].
|
|
139
|
+
|
|
140
|
+
Raises:
|
|
141
|
+
ValueError: If the images list is empty or not image objects.
|
|
142
|
+
"""
|
|
143
|
+
if not isinstance(images, list):
|
|
144
|
+
raise ValueError(
|
|
145
|
+
f"Expected a list of PIL.Image.Images, but got {type(images)}"
|
|
146
|
+
)
|
|
147
|
+
if len(images) == 0:
|
|
148
|
+
raise ValueError("Expected a non-empty list of PIL.Image.Images.")
|
|
149
|
+
if not isinstance(images[0], PIL.Image.Image):
|
|
150
|
+
raise ValueError(
|
|
151
|
+
f"Expected a list of PIL.Image.Images, but element 0 is a {type(images[0])}"
|
|
152
|
+
)
|
|
153
|
+
|
|
154
|
+
# Call the subclass implementation
|
|
155
|
+
feats = self._images_to_features(images=images)
|
|
156
|
+
if flatten:
|
|
157
|
+
# Flatten the features
|
|
158
|
+
feats = feats.reshape(feats.shape[0], -1)
|
|
159
|
+
|
|
160
|
+
return feats
|
|
161
|
+
|
|
162
|
+
@abstractmethod
|
|
163
|
+
def _images_to_features(self, images: List[PIL.Image.Image]) -> torch.Tensor:
|
|
164
|
+
"""Convert images to features.
|
|
165
|
+
|
|
166
|
+
Args:
|
|
167
|
+
images: A list of PIL.Image.Image.
|
|
168
|
+
|
|
169
|
+
Returns:
|
|
170
|
+
A torch.Tensor of shape [B, *].
|
|
171
|
+
"""
|
|
172
|
+
raise NotImplementedError
|
|
File without changes
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
from typing import List
|
|
2
|
+
|
|
3
|
+
import PIL.Image
|
|
4
|
+
import torch
|
|
5
|
+
import torchvision.transforms.v2 as v2
|
|
6
|
+
|
|
7
|
+
from enczoo.base import ImageEncoding
|
|
8
|
+
from enczoo.transforms.random_projection import RandomProjection
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
# %%
|
|
12
|
+
class Pixels(ImageEncoding):
|
|
13
|
+
"""Encode images by their resized center-crop pixels."""
|
|
14
|
+
|
|
15
|
+
def __init__(
|
|
16
|
+
self,
|
|
17
|
+
size: int = 16,
|
|
18
|
+
random_projection_dim: int | None = None,
|
|
19
|
+
random_projection_seed: int | None = None,
|
|
20
|
+
):
|
|
21
|
+
"""Initialize the pixel encoder.
|
|
22
|
+
|
|
23
|
+
Args:
|
|
24
|
+
size: Output side length in pixels.
|
|
25
|
+
random_projection_dim: Optional output dimension for projection.
|
|
26
|
+
random_projection_seed: Seed for projection weights.
|
|
27
|
+
"""
|
|
28
|
+
super().__init__()
|
|
29
|
+
|
|
30
|
+
# Register size tensor as buffer
|
|
31
|
+
self.register_buffer(
|
|
32
|
+
"size", torch.tensor(size, dtype=torch.int16, requires_grad=False)
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
# Transform
|
|
36
|
+
self.transforms = v2.Compose(
|
|
37
|
+
[
|
|
38
|
+
v2.ToImage(),
|
|
39
|
+
v2.RGB(),
|
|
40
|
+
v2.ToDtype(torch.uint8, scale=True),
|
|
41
|
+
v2.Resize(size=None, max_size=size, antialias=False),
|
|
42
|
+
v2.CenterCrop(size=size),
|
|
43
|
+
v2.ToDtype(torch.float32, scale=True),
|
|
44
|
+
]
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
# Random projection
|
|
48
|
+
if random_projection_dim is not None:
|
|
49
|
+
if random_projection_dim > (size * size * 3):
|
|
50
|
+
raise ValueError(
|
|
51
|
+
f"random_projection_dim must be less than or equal to size * size * 3={size * size * 3}!"
|
|
52
|
+
)
|
|
53
|
+
|
|
54
|
+
if random_projection_seed is None:
|
|
55
|
+
raise ValueError(
|
|
56
|
+
"random_projection_seed must be provided if random_projection_dim is not None!"
|
|
57
|
+
)
|
|
58
|
+
self.random_projection = RandomProjection(
|
|
59
|
+
seed=random_projection_seed,
|
|
60
|
+
in_features=int(size * size * 3),
|
|
61
|
+
out_features=random_projection_dim,
|
|
62
|
+
)
|
|
63
|
+
else:
|
|
64
|
+
self.random_projection = None
|
|
65
|
+
|
|
66
|
+
def _images_to_features(self, images: List[PIL.Image.Image]) -> torch.Tensor:
|
|
67
|
+
"""Convert images to pixel features.
|
|
68
|
+
|
|
69
|
+
Args:
|
|
70
|
+
images: A list of PIL.Image.Image.
|
|
71
|
+
|
|
72
|
+
Returns:
|
|
73
|
+
A torch.Tensor of shape [B, size, size, 3] or projected features.
|
|
74
|
+
"""
|
|
75
|
+
# Apply the transformations to each image
|
|
76
|
+
transformed_images = [self.transforms(image.convert("RGB")) for image in images]
|
|
77
|
+
|
|
78
|
+
# Stack the transformed images into a single tensor
|
|
79
|
+
images_tensor = torch.stack(transformed_images)
|
|
80
|
+
|
|
81
|
+
# If random projection is enabled, apply it
|
|
82
|
+
if self.random_projection is not None:
|
|
83
|
+
return self.random_projection(
|
|
84
|
+
images_tensor.reshape(images_tensor.shape[0], -1)
|
|
85
|
+
)
|
|
86
|
+
else:
|
|
87
|
+
# Rearrange from BCHW to BHWC order
|
|
88
|
+
images_tensor = images_tensor.permute(0, 2, 3, 1)
|
|
89
|
+
return images_tensor
|
|
File without changes
|
|
@@ -0,0 +1,184 @@
|
|
|
1
|
+
from abc import ABC
|
|
2
|
+
from typing import Dict, List, Tuple
|
|
3
|
+
|
|
4
|
+
import PIL.Image
|
|
5
|
+
import numpy as np
|
|
6
|
+
import torch
|
|
7
|
+
import torchvision
|
|
8
|
+
|
|
9
|
+
from enczoo.base import ImageEncoding
|
|
10
|
+
from enczoo.transforms.random_projection import RandomProjection
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class ImageNeuralNetwork(ImageEncoding, ABC):
|
|
14
|
+
"""Image encoding backed by a torch neural network."""
|
|
15
|
+
|
|
16
|
+
def __init__(
|
|
17
|
+
self,
|
|
18
|
+
image_loader: torch.nn.Module | torchvision.transforms.Compose,
|
|
19
|
+
model: torch.nn.Module,
|
|
20
|
+
layer_name: str,
|
|
21
|
+
random_projection_dim: int | None,
|
|
22
|
+
random_projection_seed: int | None,
|
|
23
|
+
):
|
|
24
|
+
"""Initialize the neural network encoder.
|
|
25
|
+
|
|
26
|
+
Args:
|
|
27
|
+
image_loader: Module that converts PIL images to model inputs.
|
|
28
|
+
model: Torch model used to compute activations.
|
|
29
|
+
layer_name: Name of the layer whose activations are returned.
|
|
30
|
+
random_projection_dim: Optional output dimension for projection.
|
|
31
|
+
random_projection_seed: Seed for projection weights.
|
|
32
|
+
|
|
33
|
+
Raises:
|
|
34
|
+
ValueError: If the layer name is not found.
|
|
35
|
+
"""
|
|
36
|
+
super().__init__()
|
|
37
|
+
|
|
38
|
+
# Ensure modules will be registered in evaluation mode
|
|
39
|
+
self.train(mode=False)
|
|
40
|
+
|
|
41
|
+
# Register buffers to ensure the model's hash is distinctive for each layer
|
|
42
|
+
self.register_buffer(
|
|
43
|
+
"layer_name", torch.tensor([ord(c) for c in layer_name], dtype=torch.int16)
|
|
44
|
+
)
|
|
45
|
+
self._layer_name = layer_name # Needed for the forward pass
|
|
46
|
+
|
|
47
|
+
def register_hook(
|
|
48
|
+
module: torch.nn.Module,
|
|
49
|
+
root_name: str,
|
|
50
|
+
activations_dict: Dict[str, torch.Tensor],
|
|
51
|
+
) -> List[str]:
|
|
52
|
+
"""Recursively register forward hooks on leaf modules.
|
|
53
|
+
|
|
54
|
+
Args:
|
|
55
|
+
module: Module whose children are walked.
|
|
56
|
+
root_name: Prefix for module names.
|
|
57
|
+
activations_dict: Dict populated with layer activations.
|
|
58
|
+
|
|
59
|
+
Returns:
|
|
60
|
+
A list of layer names in discovery order.
|
|
61
|
+
"""
|
|
62
|
+
module_names = []
|
|
63
|
+
|
|
64
|
+
nchildren = 0
|
|
65
|
+
for module_name, submodule in module.named_children():
|
|
66
|
+
if module_name != "":
|
|
67
|
+
next_root_name = (
|
|
68
|
+
root_name + "." + module_name
|
|
69
|
+
if root_name != ""
|
|
70
|
+
else module_name
|
|
71
|
+
)
|
|
72
|
+
else:
|
|
73
|
+
raise ValueError("Empty module name found in model!")
|
|
74
|
+
# Ensure module is in evaluation mode
|
|
75
|
+
submodule.train(mode=False)
|
|
76
|
+
# Recursive call:
|
|
77
|
+
submodule_names = register_hook(
|
|
78
|
+
submodule,
|
|
79
|
+
root_name=next_root_name,
|
|
80
|
+
activations_dict=activations_dict,
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
# Update the number of children
|
|
84
|
+
nchildren += 1
|
|
85
|
+
module_names.extend(submodule_names)
|
|
86
|
+
|
|
87
|
+
# Base case:
|
|
88
|
+
if nchildren == 0:
|
|
89
|
+
layer_name = root_name
|
|
90
|
+
|
|
91
|
+
if layer_name in activations_dict:
|
|
92
|
+
# Don't think this should ever happen:
|
|
93
|
+
raise Exception(
|
|
94
|
+
f"Layer name {layer_name} already exists in hidden activations! Existing keys: {self._hidden_activations.keys()}"
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
def hook_function(module: torch.nn.Module, args, output):
|
|
98
|
+
# print(f'Hook called on {layer_name}')
|
|
99
|
+
activations_dict[layer_name] = output
|
|
100
|
+
|
|
101
|
+
# Attach forward hook
|
|
102
|
+
module.register_forward_hook(hook_function)
|
|
103
|
+
|
|
104
|
+
# Base case: no children
|
|
105
|
+
module_names.append(layer_name)
|
|
106
|
+
return module_names
|
|
107
|
+
|
|
108
|
+
# Register forward hooks that will populate this dictionary with hidden activations on the forward pass:
|
|
109
|
+
self._hidden_activations: Dict[str, torch.Tensor] = {}
|
|
110
|
+
self._layer_names = register_hook(
|
|
111
|
+
model, root_name="", activations_dict=self._hidden_activations
|
|
112
|
+
)
|
|
113
|
+
|
|
114
|
+
if layer_name not in self._layer_names:
|
|
115
|
+
raise ValueError(
|
|
116
|
+
f"Layer name {layer_name} not found in model.\nAvailable layer names: {self._layer_names}"
|
|
117
|
+
)
|
|
118
|
+
|
|
119
|
+
# Populate the sizes of the layers with a forward pass
|
|
120
|
+
with torch.no_grad():
|
|
121
|
+
test_image = PIL.Image.fromarray(np.zeros((224, 224, 3), dtype=np.uint8))
|
|
122
|
+
test_image = image_loader(test_image)
|
|
123
|
+
model(test_image.unsqueeze(0))
|
|
124
|
+
|
|
125
|
+
self._layer_to_shape = {
|
|
126
|
+
layer: tuple(self._hidden_activations[layer].shape[1:])
|
|
127
|
+
for layer in self._hidden_activations
|
|
128
|
+
}
|
|
129
|
+
|
|
130
|
+
# Register modules
|
|
131
|
+
self.image_loader = image_loader
|
|
132
|
+
self.model = model
|
|
133
|
+
|
|
134
|
+
if random_projection_dim is not None:
|
|
135
|
+
if random_projection_seed is None:
|
|
136
|
+
raise ValueError(
|
|
137
|
+
"random_projection_seed must be provided if random_projection_dim is not None!"
|
|
138
|
+
)
|
|
139
|
+
self.random_projection = RandomProjection(
|
|
140
|
+
seed=random_projection_seed,
|
|
141
|
+
in_features=int(np.prod(self._layer_to_shape[layer_name])),
|
|
142
|
+
out_features=random_projection_dim,
|
|
143
|
+
)
|
|
144
|
+
else:
|
|
145
|
+
self.random_projection = None
|
|
146
|
+
|
|
147
|
+
def _images_to_features(self, images: List[PIL.Image.Image]) -> torch.Tensor:
|
|
148
|
+
"""Convert images to network activations.
|
|
149
|
+
|
|
150
|
+
Args:
|
|
151
|
+
images: A list of PIL.Image.Image.
|
|
152
|
+
|
|
153
|
+
Returns:
|
|
154
|
+
A torch.Tensor of shape [B, *].
|
|
155
|
+
"""
|
|
156
|
+
|
|
157
|
+
# Preprocess the images
|
|
158
|
+
preprocessed_images = torch.stack(
|
|
159
|
+
[self.image_loader(image) for image in images], dim=0
|
|
160
|
+
)
|
|
161
|
+
|
|
162
|
+
# Transfer to the correct device
|
|
163
|
+
preprocessed_images = preprocessed_images.to(self.device)
|
|
164
|
+
|
|
165
|
+
# Run the forward pass
|
|
166
|
+
self.model(preprocessed_images)
|
|
167
|
+
|
|
168
|
+
# Retrieve the activations for the given layer
|
|
169
|
+
f = self._hidden_activations[self._layer_name]
|
|
170
|
+
|
|
171
|
+
# Perform random projection if requested
|
|
172
|
+
if self.random_projection is not None:
|
|
173
|
+
# Flatten the features
|
|
174
|
+
f = f.reshape(f.shape[0], -1)
|
|
175
|
+
|
|
176
|
+
# Run the random projection forward
|
|
177
|
+
f = self.random_projection(f)
|
|
178
|
+
|
|
179
|
+
return f
|
|
180
|
+
|
|
181
|
+
@property
|
|
182
|
+
def layer_name_to_shape(self) -> Dict[str, Tuple[int, ...]]:
|
|
183
|
+
"""Return a mapping of layer names to activation shapes."""
|
|
184
|
+
return self._layer_to_shape
|
|
@@ -0,0 +1,163 @@
|
|
|
1
|
+
from abc import ABC, abstractmethod
|
|
2
|
+
from typing import List, Tuple
|
|
3
|
+
|
|
4
|
+
import PIL.Image
|
|
5
|
+
import torch
|
|
6
|
+
import torchvision.models
|
|
7
|
+
import torchvision.transforms.functional as F
|
|
8
|
+
|
|
9
|
+
from enczoo.neural_networks.base import ImageNeuralNetwork
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class StandardImageLoader(torch.nn.Module):
|
|
13
|
+
"""Load and normalize images for standard torchvision models.
|
|
14
|
+
|
|
15
|
+
This loader resizes directly to 224 before center-crop, preserving the
|
|
16
|
+
largest square sub-image. It also converts inputs to RGB.
|
|
17
|
+
|
|
18
|
+
Example:
|
|
19
|
+
- Original loader: resize to 256, then center-crop 224.
|
|
20
|
+
- This loader: resize to 224, then center-crop 224.
|
|
21
|
+
"""
|
|
22
|
+
|
|
23
|
+
def forward(self, img: PIL.Image.Image) -> torch.Tensor:
|
|
24
|
+
"""Convert a PIL image into a normalized tensor.
|
|
25
|
+
|
|
26
|
+
Args:
|
|
27
|
+
img: Input image.
|
|
28
|
+
|
|
29
|
+
Returns:
|
|
30
|
+
Normalized image tensor suitable for torchvision models.
|
|
31
|
+
"""
|
|
32
|
+
img = img.convert("RGB")
|
|
33
|
+
|
|
34
|
+
img_tensor = F.pil_to_tensor(pic=img)
|
|
35
|
+
img_tensor = F.resize(
|
|
36
|
+
img=img_tensor, size=[224], interpolation=F.InterpolationMode.BILINEAR
|
|
37
|
+
)
|
|
38
|
+
img_tensor = F.center_crop(img=img_tensor, output_size=[224])
|
|
39
|
+
img_tensor = F.convert_image_dtype(image=img_tensor, dtype=torch.float)
|
|
40
|
+
img_tensor = F.normalize(
|
|
41
|
+
tensor=img_tensor,
|
|
42
|
+
mean=[0.485, 0.456, 0.406],
|
|
43
|
+
std=[0.229, 0.224, 0.225],
|
|
44
|
+
)
|
|
45
|
+
return img_tensor
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
class _PretrainedNN(ImageNeuralNetwork, ABC):
|
|
49
|
+
"""Base class for pretrained torchvision encoders."""
|
|
50
|
+
|
|
51
|
+
layer_names: List[str]
|
|
52
|
+
|
|
53
|
+
def __init__(
|
|
54
|
+
self,
|
|
55
|
+
layer_name: str,
|
|
56
|
+
random_projection_dim: int | None = None,
|
|
57
|
+
random_projection_seed: int | None = None,
|
|
58
|
+
):
|
|
59
|
+
"""Initialize a pretrained encoder.
|
|
60
|
+
|
|
61
|
+
Args:
|
|
62
|
+
layer_name: Name of the layer whose activations are returned.
|
|
63
|
+
random_projection_dim: Optional output dimension for projection.
|
|
64
|
+
random_projection_seed: Seed for projection weights.
|
|
65
|
+
|
|
66
|
+
Raises:
|
|
67
|
+
ValueError: If the layer name is invalid.
|
|
68
|
+
"""
|
|
69
|
+
if layer_name not in self.layer_names:
|
|
70
|
+
raise ValueError(
|
|
71
|
+
f"Unknown layer_name: {layer_name}. Available:\n{self.layer_names}"
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
image_loader, model = self._load_modules()
|
|
75
|
+
|
|
76
|
+
# Ensure modules are in evaluation mode by default
|
|
77
|
+
image_loader.train(mode=False)
|
|
78
|
+
model.train(mode=False)
|
|
79
|
+
|
|
80
|
+
super().__init__(
|
|
81
|
+
image_loader=image_loader,
|
|
82
|
+
model=model,
|
|
83
|
+
layer_name=layer_name,
|
|
84
|
+
random_projection_dim=random_projection_dim,
|
|
85
|
+
random_projection_seed=random_projection_seed,
|
|
86
|
+
)
|
|
87
|
+
|
|
88
|
+
@abstractmethod
|
|
89
|
+
def _load_modules(self) -> Tuple[torch.nn.Module, torch.nn.Module]:
|
|
90
|
+
"""Load the image loader and model for this network.
|
|
91
|
+
|
|
92
|
+
Returns:
|
|
93
|
+
A tuple of (image_loader, model).
|
|
94
|
+
"""
|
|
95
|
+
raise NotImplementedError
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
class AlexNet(_PretrainedNN):
|
|
99
|
+
"""AlexNet encoder with named layer outputs."""
|
|
100
|
+
|
|
101
|
+
# A subset of all layers (each separated by one nonlinearity):
|
|
102
|
+
layer_names = [
|
|
103
|
+
"features.1",
|
|
104
|
+
"features.4",
|
|
105
|
+
"features.7",
|
|
106
|
+
"features.9",
|
|
107
|
+
"features.11",
|
|
108
|
+
"classifier.2",
|
|
109
|
+
"classifier.5",
|
|
110
|
+
"classifier.6",
|
|
111
|
+
]
|
|
112
|
+
|
|
113
|
+
def _load_modules(self):
|
|
114
|
+
"""Load the AlexNet image loader and model."""
|
|
115
|
+
image_loader = StandardImageLoader()
|
|
116
|
+
model = torchvision.models.alexnet(
|
|
117
|
+
weights=torchvision.models.AlexNet_Weights.IMAGENET1K_V1
|
|
118
|
+
)
|
|
119
|
+
return image_loader, model
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
class ResNet50(_PretrainedNN):
|
|
123
|
+
"""ResNet-50 encoder with named layer outputs."""
|
|
124
|
+
|
|
125
|
+
# A subset of layers (each separated by one nonlinearity, except layer4.2.relu, avgpool, and fc, which are connected by a linear layer):
|
|
126
|
+
layer_names = [
|
|
127
|
+
"relu",
|
|
128
|
+
"layer1.0.relu",
|
|
129
|
+
"layer1.1.relu",
|
|
130
|
+
"layer1.2.relu",
|
|
131
|
+
"layer2.0.relu",
|
|
132
|
+
"layer2.1.relu",
|
|
133
|
+
"layer2.2.relu",
|
|
134
|
+
"layer2.3.relu",
|
|
135
|
+
"layer3.0.relu",
|
|
136
|
+
"layer3.1.relu",
|
|
137
|
+
"layer3.2.relu",
|
|
138
|
+
"layer3.3.relu",
|
|
139
|
+
"layer3.4.relu",
|
|
140
|
+
"layer3.5.relu",
|
|
141
|
+
"layer4.0.relu",
|
|
142
|
+
"layer4.1.relu",
|
|
143
|
+
"layer4.2.relu",
|
|
144
|
+
"avgpool",
|
|
145
|
+
"fc",
|
|
146
|
+
]
|
|
147
|
+
|
|
148
|
+
def _load_modules(self):
|
|
149
|
+
"""Load the ResNet-50 image loader and model."""
|
|
150
|
+
image_loader = StandardImageLoader()
|
|
151
|
+
model = torchvision.models.resnet50(
|
|
152
|
+
weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V1
|
|
153
|
+
)
|
|
154
|
+
return image_loader, model
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
if __name__ == "__main__":
|
|
158
|
+
resnet50 = ResNet50(
|
|
159
|
+
layer_name="avgpool", random_projection_dim=None, random_projection_seed=0
|
|
160
|
+
)
|
|
161
|
+
print(resnet50.training)
|
|
162
|
+
print(resnet50.model.training)
|
|
163
|
+
print(getattr(resnet50.image_loader, "training", None))
|
|
File without changes
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
from typing import List
|
|
2
|
+
|
|
3
|
+
import PIL.Image
|
|
4
|
+
import numpy as np
|
|
5
|
+
import pytest
|
|
6
|
+
import torch
|
|
7
|
+
|
|
8
|
+
from enczoo import AlexNet
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@pytest.fixture
|
|
12
|
+
def alexnet():
|
|
13
|
+
return AlexNet(
|
|
14
|
+
layer_name=AlexNet.layer_names[-1],
|
|
15
|
+
random_projection_dim=None,
|
|
16
|
+
random_projection_seed=0,
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
@pytest.fixture
|
|
21
|
+
def images() -> List[PIL.Image.Image]:
|
|
22
|
+
np.random.seed(0)
|
|
23
|
+
img_dat1 = PIL.Image.fromarray(
|
|
24
|
+
np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
|
25
|
+
)
|
|
26
|
+
img_dat2 = PIL.Image.fromarray(
|
|
27
|
+
np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
|
28
|
+
)
|
|
29
|
+
return [img_dat1, img_dat2]
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
@pytest.fixture
|
|
33
|
+
def image_batch1(images):
|
|
34
|
+
return [images[0]]
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
@pytest.fixture
|
|
38
|
+
def image_batch2(images):
|
|
39
|
+
return images
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def test_default_is_not_training(alexnet):
|
|
43
|
+
assert not alexnet.training
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def test_alexnet_forward(alexnet, image_batch1, image_batch2):
|
|
47
|
+
# Test output is deterministic
|
|
48
|
+
result = alexnet(images=image_batch1)
|
|
49
|
+
result2 = alexnet(images=image_batch1)
|
|
50
|
+
assert torch.allclose(result, result2)
|
|
51
|
+
|
|
52
|
+
# Test output does not depend on batch size
|
|
53
|
+
result_bigger_batch = alexnet(images=image_batch2)
|
|
54
|
+
assert torch.allclose(result[0], result_bigger_batch[0], rtol=1e-3)
|
|
File without changes
|
|
File without changes
|
|
@@ -0,0 +1,73 @@
|
|
|
1
|
+
import math
|
|
2
|
+
from typing import Tuple, cast
|
|
3
|
+
|
|
4
|
+
import torch
|
|
5
|
+
import torch.nn as nn
|
|
6
|
+
import torch.nn.functional as F
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class RandomProjection(nn.Module):
|
|
10
|
+
"""Apply a fixed random projection to an input tensor."""
|
|
11
|
+
|
|
12
|
+
def __init__(
|
|
13
|
+
self,
|
|
14
|
+
in_features: int,
|
|
15
|
+
out_features: int,
|
|
16
|
+
seed: int,
|
|
17
|
+
):
|
|
18
|
+
"""Initialize the random projection.
|
|
19
|
+
|
|
20
|
+
Args:
|
|
21
|
+
in_features: Input feature dimension.
|
|
22
|
+
out_features: Output feature dimension.
|
|
23
|
+
seed: Seed for the random projection weights.
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
super().__init__()
|
|
27
|
+
self.train(mode=False)
|
|
28
|
+
|
|
29
|
+
# Register inputs as buffers; these will constitute the module's hash.
|
|
30
|
+
self.register_buffer(
|
|
31
|
+
"seed", torch.tensor(seed, dtype=torch.int64, requires_grad=False)
|
|
32
|
+
)
|
|
33
|
+
self.register_buffer(
|
|
34
|
+
"in_features",
|
|
35
|
+
torch.tensor(in_features, dtype=torch.int64, requires_grad=False),
|
|
36
|
+
)
|
|
37
|
+
self.register_buffer(
|
|
38
|
+
"out_features",
|
|
39
|
+
torch.tensor(out_features, dtype=torch.int64, requires_grad=False),
|
|
40
|
+
)
|
|
41
|
+
|
|
42
|
+
with torch.random.fork_rng():
|
|
43
|
+
# Set the weights from a standard normal distribution:
|
|
44
|
+
torch.manual_seed(seed)
|
|
45
|
+
weights = torch.randn(
|
|
46
|
+
size=(out_features, in_features),
|
|
47
|
+
dtype=torch.float32,
|
|
48
|
+
requires_grad=False,
|
|
49
|
+
) / math.sqrt(in_features * out_features)
|
|
50
|
+
|
|
51
|
+
self.register_buffer("projection_weights", weights)
|
|
52
|
+
|
|
53
|
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
54
|
+
"""Project input features with a fixed random matrix.
|
|
55
|
+
|
|
56
|
+
Args:
|
|
57
|
+
x: Input tensor of shape [B, in_features].
|
|
58
|
+
|
|
59
|
+
Returns:
|
|
60
|
+
Projected tensor of shape [B, out_features].
|
|
61
|
+
"""
|
|
62
|
+
weights = cast(torch.Tensor, self.projection_weights)
|
|
63
|
+
return F.linear(x, weights)
|
|
64
|
+
|
|
65
|
+
def __repr__(self):
|
|
66
|
+
"""Return a concise representation for debugging."""
|
|
67
|
+
return f"RandomProjection(in_features={self.in_features}, out_features={self.out_features}, seed={self.seed})"
|
|
68
|
+
|
|
69
|
+
@property
|
|
70
|
+
def output_shape(self) -> Tuple[int]:
|
|
71
|
+
"""Return the output feature shape (excluding batch dimension)."""
|
|
72
|
+
weights = cast(torch.Tensor, self.projection_weights)
|
|
73
|
+
return (weights.shape[0],)
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
import hashlib
|
|
2
|
+
import itertools
|
|
3
|
+
from typing import Iterable, Iterator, List, TypeVar
|
|
4
|
+
|
|
5
|
+
import torch
|
|
6
|
+
|
|
7
|
+
T = TypeVar("T")
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def iterate_batches(iterable: Iterable[T], batch_size: int) -> Iterator[List[T]]:
|
|
11
|
+
"""Yield items from an iterable in fixed-size batches.
|
|
12
|
+
|
|
13
|
+
Args:
|
|
14
|
+
iterable: Source iterable to batch.
|
|
15
|
+
batch_size: Number of items per batch.
|
|
16
|
+
|
|
17
|
+
Yields:
|
|
18
|
+
Lists of up to batch_size elements.
|
|
19
|
+
|
|
20
|
+
Raises:
|
|
21
|
+
ValueError: If batch_size is less than 1.
|
|
22
|
+
"""
|
|
23
|
+
if batch_size < 1:
|
|
24
|
+
raise ValueError("Batch size must be at least 1!")
|
|
25
|
+
|
|
26
|
+
iterator = iter(iterable)
|
|
27
|
+
while True:
|
|
28
|
+
batch = list(itertools.islice(iterator, batch_size))
|
|
29
|
+
if len(batch) == 0:
|
|
30
|
+
return
|
|
31
|
+
yield batch
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def hash_torch_module(module: torch.nn.Module) -> str:
|
|
35
|
+
"""Return a hash for a torch.nn.Module.
|
|
36
|
+
|
|
37
|
+
The hash depends on the module's state_dict and string representation.
|
|
38
|
+
|
|
39
|
+
Args:
|
|
40
|
+
module: Module to hash.
|
|
41
|
+
|
|
42
|
+
Returns:
|
|
43
|
+
Hex-encoded SHA-256 hash.
|
|
44
|
+
"""
|
|
45
|
+
|
|
46
|
+
sha256_hash = hashlib.sha256()
|
|
47
|
+
|
|
48
|
+
# Hash the model's state_dict:
|
|
49
|
+
state_dict = module.state_dict()
|
|
50
|
+
for key in sorted(state_dict.keys()):
|
|
51
|
+
tensor_value = state_dict[key].detach().cpu().numpy()
|
|
52
|
+
sha256_hash.update(tensor_value.tobytes())
|
|
53
|
+
|
|
54
|
+
# Hash the module's string representation:
|
|
55
|
+
module_string = str(module)
|
|
56
|
+
module_string = module_string.encode("utf-8")
|
|
57
|
+
sha256_hash.update(module_string)
|
|
58
|
+
|
|
59
|
+
# Return the combined hash:
|
|
60
|
+
return sha256_hash.hexdigest()
|