dragon-ml-toolbox 6.3.0__tar.gz → 6.4.1__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.
Potentially problematic release.
This version of dragon-ml-toolbox might be problematic. Click here for more details.
- {dragon_ml_toolbox-6.3.0/dragon_ml_toolbox.egg-info → dragon_ml_toolbox-6.4.1}/PKG-INFO +2 -1
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/README.md +1 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1/dragon_ml_toolbox.egg-info}/PKG-INFO +2 -1
- dragon_ml_toolbox-6.4.1/ml_tools/ML_inference.py +304 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ML_models.py +1 -1
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/pyproject.toml +1 -1
- dragon_ml_toolbox-6.3.0/ml_tools/ML_inference.py +0 -140
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/LICENSE +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/LICENSE-THIRD-PARTY.md +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/dragon_ml_toolbox.egg-info/SOURCES.txt +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/dragon_ml_toolbox.egg-info/dependency_links.txt +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/dragon_ml_toolbox.egg-info/requires.txt +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/dragon_ml_toolbox.egg-info/top_level.txt +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ETL_engineering.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/GUI_tools.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/MICE_imputation.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ML_callbacks.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ML_datasetmaster.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ML_evaluation.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ML_optimization.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ML_trainer.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/PSO_optimization.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/RNN_forecast.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/SQL.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/VIF_factor.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/__init__.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/_logger.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/_script_info.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/custom_logger.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/data_exploration.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ensemble_evaluation.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ensemble_inference.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/ensemble_learning.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/handle_excel.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/keys.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/optimization_tools.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/path_manager.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/ml_tools/utilities.py +0 -0
- {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/setup.cfg +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: dragon-ml-toolbox
|
|
3
|
-
Version: 6.
|
|
3
|
+
Version: 6.4.1
|
|
4
4
|
Summary: A collection of tools for data science and machine learning projects.
|
|
5
5
|
Author-email: Karl Loza <luigiloza@gmail.com>
|
|
6
6
|
License-Expression: MIT
|
|
@@ -240,6 +240,7 @@ pip install "dragon-ml-toolbox[gui-torch,plot]"
|
|
|
240
240
|
```Bash
|
|
241
241
|
custom_logger
|
|
242
242
|
GUI_tools
|
|
243
|
+
ML_models
|
|
243
244
|
ML_inference
|
|
244
245
|
path_manager
|
|
245
246
|
```
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: dragon-ml-toolbox
|
|
3
|
-
Version: 6.
|
|
3
|
+
Version: 6.4.1
|
|
4
4
|
Summary: A collection of tools for data science and machine learning projects.
|
|
5
5
|
Author-email: Karl Loza <luigiloza@gmail.com>
|
|
6
6
|
License-Expression: MIT
|
|
@@ -240,6 +240,7 @@ pip install "dragon-ml-toolbox[gui-torch,plot]"
|
|
|
240
240
|
```Bash
|
|
241
241
|
custom_logger
|
|
242
242
|
GUI_tools
|
|
243
|
+
ML_models
|
|
243
244
|
ML_inference
|
|
244
245
|
path_manager
|
|
245
246
|
```
|
|
@@ -0,0 +1,304 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
from torch import nn
|
|
3
|
+
import numpy as np
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
from typing import Union, Literal, Dict, Any, Optional
|
|
6
|
+
|
|
7
|
+
from ._script_info import _script_info
|
|
8
|
+
from ._logger import _LOGGER
|
|
9
|
+
from .path_manager import make_fullpath
|
|
10
|
+
from .keys import PyTorchInferenceKeys
|
|
11
|
+
|
|
12
|
+
__all__ = [
|
|
13
|
+
"PyTorchInferenceHandler",
|
|
14
|
+
"multi_inference_regression",
|
|
15
|
+
"multi_inference_classification"
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
class PyTorchInferenceHandler:
|
|
19
|
+
"""
|
|
20
|
+
Handles loading a PyTorch model's state dictionary and performing inference
|
|
21
|
+
for either regression or classification tasks.
|
|
22
|
+
"""
|
|
23
|
+
def __init__(self,
|
|
24
|
+
model: nn.Module,
|
|
25
|
+
state_dict: Union[str, Path],
|
|
26
|
+
task: Literal["classification", "regression"],
|
|
27
|
+
device: str = 'cpu',
|
|
28
|
+
target_id: Optional[str]=None):
|
|
29
|
+
"""
|
|
30
|
+
Initializes the handler by loading a model's state_dict.
|
|
31
|
+
|
|
32
|
+
Args:
|
|
33
|
+
model (nn.Module): An instantiated PyTorch model with the correct architecture.
|
|
34
|
+
state_dict (str | Path): The path to the saved .pth model state_dict file.
|
|
35
|
+
task (str): The type of task, 'regression' or 'classification'.
|
|
36
|
+
device (str): The device to run inference on ('cpu', 'cuda', 'mps').
|
|
37
|
+
target_id (str | None): Target name as used in the training set.
|
|
38
|
+
"""
|
|
39
|
+
self.model = model
|
|
40
|
+
self.task = task
|
|
41
|
+
self.device = self._validate_device(device)
|
|
42
|
+
self.target_id = target_id
|
|
43
|
+
|
|
44
|
+
model_p = make_fullpath(state_dict, enforce="file")
|
|
45
|
+
|
|
46
|
+
try:
|
|
47
|
+
# Load the state dictionary and apply it to the model structure
|
|
48
|
+
self.model.load_state_dict(torch.load(model_p, map_location=self.device))
|
|
49
|
+
self.model.to(self.device)
|
|
50
|
+
self.model.eval() # Set the model to evaluation mode
|
|
51
|
+
_LOGGER.info(f"✅ Model state loaded from '{model_p.name}' and set to evaluation mode.")
|
|
52
|
+
except Exception as e:
|
|
53
|
+
_LOGGER.error(f"❌ Failed to load model state from '{model_p}': {e}")
|
|
54
|
+
raise
|
|
55
|
+
|
|
56
|
+
def _validate_device(self, device: str) -> torch.device:
|
|
57
|
+
"""Validates the selected device and returns a torch.device object."""
|
|
58
|
+
device_lower = device.lower()
|
|
59
|
+
if "cuda" in device_lower and not torch.cuda.is_available():
|
|
60
|
+
_LOGGER.warning("⚠️ CUDA not available, switching to CPU.")
|
|
61
|
+
device_lower = "cpu"
|
|
62
|
+
elif device_lower == "mps" and not torch.backends.mps.is_available():
|
|
63
|
+
_LOGGER.warning("⚠️ Apple Metal Performance Shaders (MPS) not available, switching to CPU.")
|
|
64
|
+
device_lower = "cpu"
|
|
65
|
+
return torch.device(device_lower)
|
|
66
|
+
|
|
67
|
+
def _preprocess_input(self, features: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
|
|
68
|
+
"""Converts input to a torch.Tensor and moves it to the correct device."""
|
|
69
|
+
if isinstance(features, np.ndarray):
|
|
70
|
+
features = torch.from_numpy(features).float()
|
|
71
|
+
|
|
72
|
+
# Ensure tensor is on the correct device
|
|
73
|
+
return features.to(self.device)
|
|
74
|
+
|
|
75
|
+
def predict_batch(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
|
76
|
+
"""
|
|
77
|
+
Core batch prediction method. Returns results as PyTorch tensors on the model's device.
|
|
78
|
+
"""
|
|
79
|
+
if features.ndim != 2:
|
|
80
|
+
raise ValueError("Input for batch prediction must be a 2D array or tensor.")
|
|
81
|
+
|
|
82
|
+
input_tensor = self._preprocess_input(features)
|
|
83
|
+
|
|
84
|
+
with torch.no_grad():
|
|
85
|
+
# Output tensor remains on the model's device (e.g., 'mps' or 'cuda')
|
|
86
|
+
output = self.model(input_tensor)
|
|
87
|
+
|
|
88
|
+
if self.task == "classification":
|
|
89
|
+
probs = nn.functional.softmax(output, dim=1)
|
|
90
|
+
labels = torch.argmax(probs, dim=1)
|
|
91
|
+
return {
|
|
92
|
+
PyTorchInferenceKeys.LABELS: labels,
|
|
93
|
+
PyTorchInferenceKeys.PROBABILITIES: probs
|
|
94
|
+
}
|
|
95
|
+
else: # regression
|
|
96
|
+
return {PyTorchInferenceKeys.PREDICTIONS: output}
|
|
97
|
+
|
|
98
|
+
def predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
|
99
|
+
"""
|
|
100
|
+
Core single-sample prediction. Returns results as PyTorch tensors on the model's device.
|
|
101
|
+
"""
|
|
102
|
+
if features.ndim == 1:
|
|
103
|
+
features = features.reshape(1, -1)
|
|
104
|
+
|
|
105
|
+
if features.shape[0] != 1:
|
|
106
|
+
raise ValueError("The predict() method is for a single sample. Use predict_batch() for multiple samples.")
|
|
107
|
+
|
|
108
|
+
batch_results = self.predict_batch(features)
|
|
109
|
+
|
|
110
|
+
single_results = {key: value[0] for key, value in batch_results.items()}
|
|
111
|
+
return single_results
|
|
112
|
+
|
|
113
|
+
# --- NumPy Convenience Wrappers (on CPU) ---
|
|
114
|
+
|
|
115
|
+
def predict_batch_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, np.ndarray]:
|
|
116
|
+
"""
|
|
117
|
+
Convenience wrapper for predict_batch that returns NumPy arrays.
|
|
118
|
+
"""
|
|
119
|
+
tensor_results = self.predict_batch(features)
|
|
120
|
+
# Move tensor to CPU before converting to NumPy
|
|
121
|
+
numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
|
|
122
|
+
return numpy_results
|
|
123
|
+
|
|
124
|
+
def predict_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
|
|
125
|
+
"""
|
|
126
|
+
Convenience wrapper for predict that returns NumPy arrays or scalars.
|
|
127
|
+
"""
|
|
128
|
+
tensor_results = self.predict(features)
|
|
129
|
+
|
|
130
|
+
if self.task == "regression":
|
|
131
|
+
# .item() implicitly moves to CPU
|
|
132
|
+
return {PyTorchInferenceKeys.PREDICTIONS: tensor_results[PyTorchInferenceKeys.PREDICTIONS].item()}
|
|
133
|
+
else: # classification
|
|
134
|
+
return {
|
|
135
|
+
PyTorchInferenceKeys.LABELS: tensor_results[PyTorchInferenceKeys.LABELS].item(),
|
|
136
|
+
# Move tensor to CPU before converting to NumPy
|
|
137
|
+
PyTorchInferenceKeys.PROBABILITIES: tensor_results[PyTorchInferenceKeys.PROBABILITIES].cpu().numpy()
|
|
138
|
+
}
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
def multi_inference_regression(handlers: list[PyTorchInferenceHandler],
|
|
142
|
+
feature_vector: Union[np.ndarray, torch.Tensor],
|
|
143
|
+
output: Literal["numpy","torch"]="numpy") -> dict[str,Any]:
|
|
144
|
+
"""
|
|
145
|
+
Performs regression inference using multiple models on a single feature vector.
|
|
146
|
+
|
|
147
|
+
This function iterates through a list of PyTorchInferenceHandler objects,
|
|
148
|
+
each configured for a different regression target. It runs a prediction for
|
|
149
|
+
each handler using the same input feature vector and returns the results
|
|
150
|
+
in a dictionary.
|
|
151
|
+
|
|
152
|
+
The function adapts its behavior based on the input dimensions:
|
|
153
|
+
- 1D input: Returns a dictionary mapping target ID to a single value.
|
|
154
|
+
- 2D input: Returns a dictionary mapping target ID to a list of values.
|
|
155
|
+
|
|
156
|
+
Args:
|
|
157
|
+
handlers (list[PyTorchInferenceHandler]): A list of initialized inference
|
|
158
|
+
handlers. Each handler must have a unique `target_id` and be configured with `task="regression"`.
|
|
159
|
+
feature_vector (Union[np.ndarray, torch.Tensor]): An input sample (1D) or a batch of samples (2D) to be fed into each regression model.
|
|
160
|
+
output (Literal["numpy", "torch"], optional): The desired format for the output predictions.
|
|
161
|
+
- "numpy": Returns predictions as Python scalars or NumPy arrays.
|
|
162
|
+
- "torch": Returns predictions as PyTorch tensors.
|
|
163
|
+
|
|
164
|
+
Returns:
|
|
165
|
+
(dict[str, Any]): A dictionary mapping each handler's `target_id` to its
|
|
166
|
+
predicted regression values.
|
|
167
|
+
|
|
168
|
+
Raises:
|
|
169
|
+
AttributeError: If any handler in the list is missing a `target_id`.
|
|
170
|
+
ValueError: If any handler's `task` is not 'regression' or if the input `feature_vector` is not 1D or 2D.
|
|
171
|
+
"""
|
|
172
|
+
# check batch dimension
|
|
173
|
+
is_single_sample = feature_vector.ndim == 1
|
|
174
|
+
|
|
175
|
+
# Reshape a 1D vector to a 2D batch of one for uniform processing.
|
|
176
|
+
if is_single_sample:
|
|
177
|
+
feature_vector = feature_vector.reshape(1, -1)
|
|
178
|
+
|
|
179
|
+
# Validate that the input is a 2D tensor.
|
|
180
|
+
if feature_vector.ndim != 2:
|
|
181
|
+
raise ValueError("Input feature_vector must be a 1D or 2D array/tensor.")
|
|
182
|
+
|
|
183
|
+
results: dict[str,Any] = dict()
|
|
184
|
+
for handler in handlers:
|
|
185
|
+
# validation
|
|
186
|
+
if handler.target_id is None:
|
|
187
|
+
raise AttributeError("All inference handlers must have a 'target_id' attribute.")
|
|
188
|
+
if handler.task != "regression":
|
|
189
|
+
raise ValueError(
|
|
190
|
+
f"Invalid task type: The handler for target_id '{handler.target_id}' "
|
|
191
|
+
f"is for '{handler.task}', but only 'regression' tasks are supported."
|
|
192
|
+
)
|
|
193
|
+
|
|
194
|
+
# inference
|
|
195
|
+
if output == "numpy":
|
|
196
|
+
# This path returns NumPy arrays or standard Python scalars
|
|
197
|
+
numpy_result = handler.predict_batch_numpy(feature_vector)[PyTorchInferenceKeys.PREDICTIONS]
|
|
198
|
+
if is_single_sample:
|
|
199
|
+
# For a single sample, convert the 1-element array to a Python scalar
|
|
200
|
+
results[handler.target_id] = numpy_result.item()
|
|
201
|
+
else:
|
|
202
|
+
# For a batch, return the full NumPy array of predictions
|
|
203
|
+
results[handler.target_id] = numpy_result
|
|
204
|
+
|
|
205
|
+
else: # output == "torch"
|
|
206
|
+
# This path returns PyTorch tensors on the model's device
|
|
207
|
+
torch_result = handler.predict_batch(feature_vector)[PyTorchInferenceKeys.PREDICTIONS]
|
|
208
|
+
if is_single_sample:
|
|
209
|
+
# For a single sample, return the 0-dim tensor
|
|
210
|
+
results[handler.target_id] = torch_result[0]
|
|
211
|
+
else:
|
|
212
|
+
# For a batch, return the full tensor of predictions
|
|
213
|
+
results[handler.target_id] = torch_result
|
|
214
|
+
|
|
215
|
+
return results
|
|
216
|
+
|
|
217
|
+
|
|
218
|
+
def multi_inference_classification(
|
|
219
|
+
handlers: list[PyTorchInferenceHandler],
|
|
220
|
+
feature_vector: Union[np.ndarray, torch.Tensor],
|
|
221
|
+
output: Literal["numpy","torch"]="numpy"
|
|
222
|
+
) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
223
|
+
"""
|
|
224
|
+
Performs classification inference on a single sample or a batch.
|
|
225
|
+
|
|
226
|
+
This function iterates through a list of PyTorchInferenceHandler objects,
|
|
227
|
+
each configured for a different classification target. It returns two
|
|
228
|
+
dictionaries: one for the predicted labels and one for the probabilities.
|
|
229
|
+
|
|
230
|
+
The function adapts its behavior based on the input dimensions:
|
|
231
|
+
- 1D input: The dictionaries map target ID to a single label and a single probability array.
|
|
232
|
+
- 2D input: The dictionaries map target ID to an array of labels and an array of probability arrays.
|
|
233
|
+
|
|
234
|
+
Args:
|
|
235
|
+
handlers (list[PyTorchInferenceHandler]): A list of initialized inference handlers. Each must have a unique `target_id` and be configured
|
|
236
|
+
with `task="classification"`.
|
|
237
|
+
feature_vector (Union[np.ndarray, torch.Tensor]): An input sample (1D)
|
|
238
|
+
or a batch of samples (2D) for prediction.
|
|
239
|
+
output (Literal["numpy", "torch"], optional): The desired format for the
|
|
240
|
+
output predictions.
|
|
241
|
+
|
|
242
|
+
Returns:
|
|
243
|
+
(tuple[dict[str, Any], dict[str, Any]]): A tuple containing two dictionaries:
|
|
244
|
+
1. A dictionary mapping `target_id` to the predicted label(s).
|
|
245
|
+
2. A dictionary mapping `target_id` to the prediction probabilities.
|
|
246
|
+
|
|
247
|
+
Raises:
|
|
248
|
+
AttributeError: If any handler in the list is missing a `target_id`.
|
|
249
|
+
ValueError: If any handler's `task` is not 'classification' or if the input `feature_vector` is not 1D or 2D.
|
|
250
|
+
"""
|
|
251
|
+
# Store if the original input was a single sample
|
|
252
|
+
is_single_sample = feature_vector.ndim == 1
|
|
253
|
+
|
|
254
|
+
# Reshape a 1D vector to a 2D batch of one for uniform processing
|
|
255
|
+
if is_single_sample:
|
|
256
|
+
feature_vector = feature_vector.reshape(1, -1)
|
|
257
|
+
|
|
258
|
+
if feature_vector.ndim != 2:
|
|
259
|
+
raise ValueError("Input feature_vector must be a 1D or 2D array/tensor.")
|
|
260
|
+
|
|
261
|
+
# Initialize two dictionaries for results
|
|
262
|
+
labels_results: dict[str, Any] = dict()
|
|
263
|
+
probs_results: dict[str, Any] = dict()
|
|
264
|
+
|
|
265
|
+
for handler in handlers:
|
|
266
|
+
# Validation
|
|
267
|
+
if handler.target_id is None:
|
|
268
|
+
raise AttributeError("All inference handlers must have a 'target_id' attribute.")
|
|
269
|
+
if handler.task != "classification":
|
|
270
|
+
raise ValueError(
|
|
271
|
+
f"Invalid task type: The handler for target_id '{handler.target_id}' "
|
|
272
|
+
f"is for '{handler.task}', but this function only supports 'classification'."
|
|
273
|
+
)
|
|
274
|
+
|
|
275
|
+
# Inference
|
|
276
|
+
if output == "numpy":
|
|
277
|
+
# predict_batch_numpy returns a dict of NumPy arrays
|
|
278
|
+
result = handler.predict_batch_numpy(feature_vector)
|
|
279
|
+
else: # torch
|
|
280
|
+
# predict_batch returns a dict of Torch tensors
|
|
281
|
+
result = handler.predict_batch(feature_vector)
|
|
282
|
+
|
|
283
|
+
labels = result[PyTorchInferenceKeys.LABELS]
|
|
284
|
+
probabilities = result[PyTorchInferenceKeys.PROBABILITIES]
|
|
285
|
+
|
|
286
|
+
if is_single_sample:
|
|
287
|
+
# For "numpy", convert the single label to a Python int scalar.
|
|
288
|
+
# For "torch", get the 0-dim tensor label.
|
|
289
|
+
if output == "numpy":
|
|
290
|
+
labels_results[handler.target_id] = labels.item()
|
|
291
|
+
else: # torch
|
|
292
|
+
labels_results[handler.target_id] = labels[0]
|
|
293
|
+
|
|
294
|
+
# The probabilities are an array/tensor of values
|
|
295
|
+
probs_results[handler.target_id] = probabilities[0]
|
|
296
|
+
else:
|
|
297
|
+
labels_results[handler.target_id] = labels
|
|
298
|
+
probs_results[handler.target_id] = probabilities
|
|
299
|
+
|
|
300
|
+
return labels_results, probs_results
|
|
301
|
+
|
|
302
|
+
|
|
303
|
+
def info():
|
|
304
|
+
_script_info(__all__)
|
|
@@ -223,7 +223,7 @@ def save_architecture(model: nn.Module, directory: Union[str, Path], verbose: bo
|
|
|
223
223
|
json.dump(config, f, indent=4)
|
|
224
224
|
|
|
225
225
|
if verbose:
|
|
226
|
-
_LOGGER.info(f"✅ Architecture for '{model.__class__.__name__}' saved to '{path_dir}'")
|
|
226
|
+
_LOGGER.info(f"✅ Architecture for '{model.__class__.__name__}' saved to '{path_dir.name}'")
|
|
227
227
|
|
|
228
228
|
|
|
229
229
|
def load_architecture(filepath: Union[str, Path], expected_model_class: type, verbose: bool=True) -> nn.Module:
|
|
@@ -1,140 +0,0 @@
|
|
|
1
|
-
import torch
|
|
2
|
-
from torch import nn
|
|
3
|
-
import numpy as np
|
|
4
|
-
from pathlib import Path
|
|
5
|
-
from typing import Union, Literal, Dict, Any, Optional
|
|
6
|
-
|
|
7
|
-
from ._script_info import _script_info
|
|
8
|
-
from ._logger import _LOGGER
|
|
9
|
-
from .path_manager import make_fullpath
|
|
10
|
-
from .keys import PyTorchInferenceKeys
|
|
11
|
-
|
|
12
|
-
__all__ = [
|
|
13
|
-
"PyTorchInferenceHandler"
|
|
14
|
-
]
|
|
15
|
-
|
|
16
|
-
class PyTorchInferenceHandler:
|
|
17
|
-
"""
|
|
18
|
-
Handles loading a PyTorch model's state dictionary and performing inference
|
|
19
|
-
for either regression or classification tasks.
|
|
20
|
-
"""
|
|
21
|
-
def __init__(self,
|
|
22
|
-
model: nn.Module,
|
|
23
|
-
state_dict: Union[str, Path],
|
|
24
|
-
task: Literal["classification", "regression"],
|
|
25
|
-
device: str = 'cpu',
|
|
26
|
-
target_id: Optional[str]=None):
|
|
27
|
-
"""
|
|
28
|
-
Initializes the handler by loading a model's state_dict.
|
|
29
|
-
|
|
30
|
-
Args:
|
|
31
|
-
model (nn.Module): An instantiated PyTorch model with the correct architecture.
|
|
32
|
-
state_dict (str | Path): The path to the saved .pth model state_dict file.
|
|
33
|
-
task (str): The type of task, 'regression' or 'classification'.
|
|
34
|
-
device (str): The device to run inference on ('cpu', 'cuda', 'mps').
|
|
35
|
-
target_id (str | None): Target name as used in the training set.
|
|
36
|
-
"""
|
|
37
|
-
self.model = model
|
|
38
|
-
self.task = task
|
|
39
|
-
self.device = self._validate_device(device)
|
|
40
|
-
self.target_id = target_id
|
|
41
|
-
|
|
42
|
-
model_p = make_fullpath(state_dict, enforce="file")
|
|
43
|
-
|
|
44
|
-
try:
|
|
45
|
-
# Load the state dictionary and apply it to the model structure
|
|
46
|
-
self.model.load_state_dict(torch.load(model_p, map_location=self.device))
|
|
47
|
-
self.model.to(self.device)
|
|
48
|
-
self.model.eval() # Set the model to evaluation mode
|
|
49
|
-
_LOGGER.info(f"✅ Model state loaded from '{model_p.name}' and set to evaluation mode.")
|
|
50
|
-
except Exception as e:
|
|
51
|
-
_LOGGER.error(f"❌ Failed to load model state from '{model_p}': {e}")
|
|
52
|
-
raise
|
|
53
|
-
|
|
54
|
-
def _validate_device(self, device: str) -> torch.device:
|
|
55
|
-
"""Validates the selected device and returns a torch.device object."""
|
|
56
|
-
device_lower = device.lower()
|
|
57
|
-
if "cuda" in device_lower and not torch.cuda.is_available():
|
|
58
|
-
_LOGGER.warning("⚠️ CUDA not available, switching to CPU.")
|
|
59
|
-
device_lower = "cpu"
|
|
60
|
-
elif device_lower == "mps" and not torch.backends.mps.is_available():
|
|
61
|
-
_LOGGER.warning("⚠️ Apple Metal Performance Shaders (MPS) not available, switching to CPU.")
|
|
62
|
-
device_lower = "cpu"
|
|
63
|
-
return torch.device(device_lower)
|
|
64
|
-
|
|
65
|
-
def _preprocess_input(self, features: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
|
|
66
|
-
"""Converts input to a torch.Tensor and moves it to the correct device."""
|
|
67
|
-
if isinstance(features, np.ndarray):
|
|
68
|
-
features = torch.from_numpy(features).float()
|
|
69
|
-
|
|
70
|
-
# Ensure tensor is on the correct device
|
|
71
|
-
return features.to(self.device)
|
|
72
|
-
|
|
73
|
-
def predict_batch(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
|
74
|
-
"""
|
|
75
|
-
Core batch prediction method. Returns results as PyTorch tensors on the model's device.
|
|
76
|
-
"""
|
|
77
|
-
if features.ndim != 2:
|
|
78
|
-
raise ValueError("Input for batch prediction must be a 2D array or tensor.")
|
|
79
|
-
|
|
80
|
-
input_tensor = self._preprocess_input(features)
|
|
81
|
-
|
|
82
|
-
with torch.no_grad():
|
|
83
|
-
# Output tensor remains on the model's device (e.g., 'mps' or 'cuda')
|
|
84
|
-
output = self.model(input_tensor)
|
|
85
|
-
|
|
86
|
-
if self.task == "classification":
|
|
87
|
-
probs = nn.functional.softmax(output, dim=1)
|
|
88
|
-
labels = torch.argmax(probs, dim=1)
|
|
89
|
-
return {
|
|
90
|
-
PyTorchInferenceKeys.LABELS: labels,
|
|
91
|
-
PyTorchInferenceKeys.PROBABILITIES: probs
|
|
92
|
-
}
|
|
93
|
-
else: # regression
|
|
94
|
-
return {PyTorchInferenceKeys.PREDICTIONS: output}
|
|
95
|
-
|
|
96
|
-
def predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
|
97
|
-
"""
|
|
98
|
-
Core single-sample prediction. Returns results as PyTorch tensors on the model's device.
|
|
99
|
-
"""
|
|
100
|
-
if features.ndim == 1:
|
|
101
|
-
features = features.reshape(1, -1)
|
|
102
|
-
|
|
103
|
-
if features.shape[0] != 1:
|
|
104
|
-
raise ValueError("The predict() method is for a single sample. Use predict_batch() for multiple samples.")
|
|
105
|
-
|
|
106
|
-
batch_results = self.predict_batch(features)
|
|
107
|
-
|
|
108
|
-
single_results = {key: value[0] for key, value in batch_results.items()}
|
|
109
|
-
return single_results
|
|
110
|
-
|
|
111
|
-
# --- NumPy Convenience Wrappers (on CPU) ---
|
|
112
|
-
|
|
113
|
-
def predict_batch_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, np.ndarray]:
|
|
114
|
-
"""
|
|
115
|
-
Convenience wrapper for predict_batch that returns NumPy arrays.
|
|
116
|
-
"""
|
|
117
|
-
tensor_results = self.predict_batch(features)
|
|
118
|
-
# Move tensor to CPU before converting to NumPy
|
|
119
|
-
numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
|
|
120
|
-
return numpy_results
|
|
121
|
-
|
|
122
|
-
def predict_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
|
|
123
|
-
"""
|
|
124
|
-
Convenience wrapper for predict that returns NumPy arrays or scalars.
|
|
125
|
-
"""
|
|
126
|
-
tensor_results = self.predict(features)
|
|
127
|
-
|
|
128
|
-
if self.task == "regression":
|
|
129
|
-
# .item() implicitly moves to CPU
|
|
130
|
-
return {PyTorchInferenceKeys.PREDICTIONS: tensor_results[PyTorchInferenceKeys.PREDICTIONS].item()}
|
|
131
|
-
else: # classification
|
|
132
|
-
return {
|
|
133
|
-
PyTorchInferenceKeys.LABELS: tensor_results[PyTorchInferenceKeys.LABELS].item(),
|
|
134
|
-
# ✅ Move tensor to CPU before converting to NumPy
|
|
135
|
-
PyTorchInferenceKeys.PROBABILITIES: tensor_results[PyTorchInferenceKeys.PROBABILITIES].cpu().numpy()
|
|
136
|
-
}
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
def info():
|
|
140
|
-
_script_info(__all__)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/dragon_ml_toolbox.egg-info/dependency_links.txt
RENAMED
|
File without changes
|
|
File without changes
|
{dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.1}/dragon_ml_toolbox.egg-info/top_level.txt
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|