src.models

Neural network model definitions for Hepatitis C classification.

This module contains only the model architecture definitions, following the separation of concerns principle. Training logic is in train.py and data handling is in data.py.

  1"""
  2Neural network model definitions for Hepatitis C classification.
  3
  4This module contains only the model architecture definitions,
  5following the separation of concerns principle. Training logic is in train.py
  6and data handling is in data.py.
  7"""
  8
  9from __future__ import annotations
 10import torch
 11import torch.nn as nn
 12import torch.optim as optim
 13from torch.utils.data import DataLoader
 14import numpy as np
 15from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
 16from sklearn.base import BaseEstimator, ClassifierMixin
 17import tempfile
 18import os
 19
 20class ResidualBlock(nn.Module):
 21    """
 22    Residual block with layer normalization and dropout.
 23
 24    Parameters
 25    -----------
 26    size : int
 27        Size of the input and output features.
 28    dropout_rate : float
 29        Dropout rate for regularization.
 30
 31    Attributes
 32    -----------
 33    block : nn.Sequential
 34        Sequential container for the residual block layers.
 35        
 36    Examples
 37    ---------
 38    >>> block = ResidualBlock(size=128, dropout_rate=0.3)
 39    >>> input_tensor = torch.randn(32, 128)
 40    >>> output_tensor = block(input_tensor)
 41    """
 42    def __init__(self, size: int, dropout_rate: float = 0.3):
 43        super(ResidualBlock, self).__init__()
 44        self.block = nn.Sequential(
 45            nn.LayerNorm(size),
 46            nn.Linear(size, size),
 47            nn.ReLU(),
 48            nn.Dropout(dropout_rate),
 49            nn.LayerNorm(size),
 50            nn.Linear(size, size),
 51            nn.Dropout(dropout_rate)
 52        )
 53    
 54    def forward(self, x: torch.Tensor) -> torch.Tensor:
 55        return x + self.block(x)
 56
 57class HepatitisNet(nn.Module):
 58    """
 59    Neural Network for Hepatitis C classification with residual connections.
 60
 61    Parameters
 62    -----------
 63    input_size : int
 64        Number of input features.
 65    hidden_sizes : list of int
 66        List of hidden layer sizes.
 67    num_classes : int
 68        Number of output classes.
 69    dropout_rate : float
 70        Dropout rate for regularization.
 71    num_residual_blocks : int
 72        Number of residual blocks to use.
 73
 74    Attributes
 75    -----------
 76    layers : nn.ModuleList
 77        List of network layers including residual blocks.
 78    input_size : int
 79        Number of input features.
 80    num_classes : int
 81        Number of output classes.
 82    """
 83
 84
 85    def __init__(self, input_size: int = 12, hidden_sizes: list = [128, 64, 32], 
 86                 num_classes: int = 2, dropout_rate: float = 0.3, num_residual_blocks: int = 2):
 87        super(HepatitisNet, self).__init__()
 88        
 89        self.input_size = input_size
 90        self.num_classes = num_classes
 91
 92        # Build network architecture
 93        layers = nn.ModuleList()
 94        
 95        # Input projection
 96        layers.append(nn.Linear(input_size, hidden_sizes[0]))
 97        layers.append(nn.LayerNorm(hidden_sizes[0]))
 98        layers.append(nn.ReLU())
 99        layers.append(nn.Dropout(dropout_rate))
100        
101        # Add residual blocks at each hidden layer
102        for i in range(len(hidden_sizes) - 1):
103            # Add residual blocks
104            for _ in range(num_residual_blocks):
105                layers.append(ResidualBlock(hidden_sizes[i], dropout_rate))
106            
107            # Project to next hidden size
108            layers.append(nn.Linear(hidden_sizes[i], hidden_sizes[i + 1]))
109            layers.append(nn.LayerNorm(hidden_sizes[i + 1]))
110            layers.append(nn.ReLU())
111            layers.append(nn.Dropout(dropout_rate))
112        
113        # Add final residual blocks
114        for _ in range(num_residual_blocks):
115            layers.append(ResidualBlock(hidden_sizes[-1], dropout_rate))
116        
117        # Output projection
118        layers.append(nn.Linear(hidden_sizes[-1], num_classes))
119        
120        self.layers = layers
121        self._initialize_weights()
122    
123    def _initialize_weights(self):
124        for module in self.modules():
125            if isinstance(module, nn.Linear):
126                nn.init.xavier_uniform_(module.weight)
127                if module.bias is not None:
128                    nn.init.constant_(module.bias, 0)
129            elif isinstance(module, nn.BatchNorm1d):
130                nn.init.constant_(module.weight, 1)
131                nn.init.constant_(module.bias, 0)
132    
133    def forward(self, x: torch.Tensor) -> torch.Tensor:
134        for layer in self.layers:
135            x = layer(x)
136        return x
137
138
139def evaluate_model(model: nn.Module, test_loader: DataLoader, device: str = 'cpu') -> tuple[np.ndarray, np.ndarray, np.ndarray]:
140    """
141    Evaluate the model on the test dataset.
142    
143    Parameters
144    -----------
145    model : nn.Module
146        The trained model to be evaluated.
147    test_loader : DataLoader
148        DataLoader for the test dataset.
149    device : str
150        Device to run the evaluation on (default: 'cpu').
151
152    Returns
153    -----------
154    tuple[np.ndarray, np.ndarray, np.ndarray]
155        - y_true: Ground truth labels.
156        - y_pred: Predicted labels.
157        - y_probs: Predicted probabilities.
158
159    Examples
160    ---------
161    >>> y_true, y_pred, y_probs = evaluate_model(model, test_loader, device='cuda')
162    """
163    model.eval()
164    y_true = []
165    y_pred = []
166    y_probs = []
167    
168    with torch.no_grad():
169        for data, target in test_loader:
170            data, target = data.to(device), target.to(device)
171            output = model(data)
172            probs = torch.softmax(output, dim=1)
173            pred = output.argmax(dim=1)
174            
175            y_true.extend(target.cpu().numpy())
176            y_pred.extend(pred.cpu().numpy())
177            y_probs.extend(probs.cpu().numpy())
178    
179    return np.array(y_true), np.array(y_pred), np.array(y_probs)
180
181def save_model(model: nn.Module, filepath: str, additional_info: dict = None, demo: bool = False) -> None:
182    """
183    Save the model to a file.
184
185    Parameters
186    -----------
187    model : nn.Module
188        The model to be saved.
189    filepath : str
190        Path to the file where the model will be saved.
191    additional_info : dict, optional
192        Any additional information to save with the model (e.g., training parameters).
193    demo : bool, optional
194        Whether the model is being saved in a temp location demo mode (default: False).
195
196    Returns
197    -----------
198    str
199        The path to the saved model file.
200
201    Examples
202    ---------
203    >>> save_model(model, 'models/hepatitis_model.pth', {'input_size': 12, 'num_classes': 2})
204    """
205
206    if demo:
207        filepath = os.path.join(tempfile.gettempdir(), 'hepatitis_model.pth')
208    torch.save({
209        'model_state_dict': model.state_dict(),
210        'model_class': model.__class__.__name__,
211        'additional_info': additional_info
212    }, filepath, _use_new_zipfile_serialization=False)
213    print(f"Model saved to: {filepath}")
214    return filepath
215
216def load_model(filepath: str, model_class: type[HepatitisNet] = HepatitisNet, input_size: int = 12) -> tuple[nn.Module, dict]:
217    """
218    Load a model from a file.
219    
220    Parameters
221    -----------
222    filepath : str
223        Path to the file from which the model will be loaded.
224    model_class : type
225        The class of the model to be loaded (default: HepatitisNet).
226    input_size : int
227        Number of input features (default: 12).
228
229    Returns
230    -----------
231    tuple[nn.Module, dict]
232        - model: The loaded model.
233        - additional_info: Any additional information saved with the model.
234    
235    Examples
236    ---------
237    >>> model, info = load_model('models/hepatitis_model.pth')
238    >>> print(info)
239    """
240    checkpoint = torch.load(filepath, map_location='cpu', weights_only=False)
241    
242    model = model_class(input_size=input_size)
243    model.load_state_dict(checkpoint['model_state_dict'])
244    model.eval()
245    
246    return model, checkpoint.get('additional_info', None)
247
248class TorchWrapper(BaseEstimator, ClassifierMixin):
249    """
250    A wrapper to make PyTorch models compatible with scikit-learn.
251    
252    Parameters
253    -----------
254    model : HepatitisNet
255        The PyTorch model instance.
256    device : str
257        Device to run inference on ('cpu' or 'cuda').
258    classes : array-like
259        Class labels for the classifier.
260
261    Attributes
262    -----------
263    model : HepatitisNet
264        The PyTorch model instance.
265    device : str
266        Device to run inference on ('cpu' or 'cuda').
267    classes_ : array-like
268        Class labels for the classifier.
269
270    Examples
271    ---------
272    >>> model, _ = load_model('models/hepatitis_model.pth')
273    >>> wrapper = TorchWrapper(model, device='cuda', classes=[0, 1])
274    >>> CalibrationDisplay.from_estimator(wrapper, X_test, y_test, n_bins=10)
275    """
276
277    def __init__(self, model: type[HepatitisNet], device: str, classes: any):
278        self.model = model
279        self.device = device
280        self.classes_ = classes
281    def fit(self, X: np.ndarray, y: np.ndarray) -> TorchWrapper:
282        return self
283    
284    def predict_proba(self, X: np.ndarray) -> np.ndarray:
285        X_tensor = torch.FloatTensor(X).to(self.device)
286        with torch.no_grad():
287            outputs = self.model(X_tensor)
288            probs = torch.softmax(outputs, dim=1).cpu().numpy()
289        return probs
290    
291    def predict(self, X: np.ndarray) -> np.ndarray:
292        return np.argmax(self.predict_proba(X), axis=1)
class ResidualBlock(torch.nn.modules.module.Module): 
21class ResidualBlock(nn.Module):
22    """
23    Residual block with layer normalization and dropout.
24
25    Parameters
26    -----------
27    size : int
28        Size of the input and output features.
29    dropout_rate : float
30        Dropout rate for regularization.
31
32    Attributes
33    -----------
34    block : nn.Sequential
35        Sequential container for the residual block layers.
36        
37    Examples
38    ---------
39    >>> block = ResidualBlock(size=128, dropout_rate=0.3)
40    >>> input_tensor = torch.randn(32, 128)
41    >>> output_tensor = block(input_tensor)
42    """
43    def __init__(self, size: int, dropout_rate: float = 0.3):
44        super(ResidualBlock, self).__init__()
45        self.block = nn.Sequential(
46            nn.LayerNorm(size),
47            nn.Linear(size, size),
48            nn.ReLU(),
49            nn.Dropout(dropout_rate),
50            nn.LayerNorm(size),
51            nn.Linear(size, size),
52            nn.Dropout(dropout_rate)
53        )
54    
55    def forward(self, x: torch.Tensor) -> torch.Tensor:
56        return x + self.block(x)

Residual block with layer normalization and dropout.

Parameters
  • size (int): Size of the input and output features.
  • dropout_rate (float): Dropout rate for regularization.
Attributes
  • block (nn.Sequential): Sequential container for the residual block layers.
Examples
>>> block = ResidualBlock(size=128, dropout_rate=0.3)
>>> input_tensor = torch.randn(32, 128)
>>> output_tensor = block(input_tensor)
ResidualBlock(size: int, dropout_rate: float = 0.3) 
43    def __init__(self, size: int, dropout_rate: float = 0.3):
44        super(ResidualBlock, self).__init__()
45        self.block = nn.Sequential(
46            nn.LayerNorm(size),
47            nn.Linear(size, size),
48            nn.ReLU(),
49            nn.Dropout(dropout_rate),
50            nn.LayerNorm(size),
51            nn.Linear(size, size),
52            nn.Dropout(dropout_rate)
53        )

Initialize internal Module state, shared by both nn.Module and ScriptModule.

block
def forward(self, x: torch.Tensor) -> torch.Tensor: 
55    def forward(self, x: torch.Tensor) -> torch.Tensor:
56        return x + self.block(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HepatitisNet(torch.nn.modules.module.Module): 
 58class HepatitisNet(nn.Module):
 59    """
 60    Neural Network for Hepatitis C classification with residual connections.
 61
 62    Parameters
 63    -----------
 64    input_size : int
 65        Number of input features.
 66    hidden_sizes : list of int
 67        List of hidden layer sizes.
 68    num_classes : int
 69        Number of output classes.
 70    dropout_rate : float
 71        Dropout rate for regularization.
 72    num_residual_blocks : int
 73        Number of residual blocks to use.
 74
 75    Attributes
 76    -----------
 77    layers : nn.ModuleList
 78        List of network layers including residual blocks.
 79    input_size : int
 80        Number of input features.
 81    num_classes : int
 82        Number of output classes.
 83    """
 84
 85
 86    def __init__(self, input_size: int = 12, hidden_sizes: list = [128, 64, 32], 
 87                 num_classes: int = 2, dropout_rate: float = 0.3, num_residual_blocks: int = 2):
 88        super(HepatitisNet, self).__init__()
 89        
 90        self.input_size = input_size
 91        self.num_classes = num_classes
 92
 93        # Build network architecture
 94        layers = nn.ModuleList()
 95        
 96        # Input projection
 97        layers.append(nn.Linear(input_size, hidden_sizes[0]))
 98        layers.append(nn.LayerNorm(hidden_sizes[0]))
 99        layers.append(nn.ReLU())
100        layers.append(nn.Dropout(dropout_rate))
101        
102        # Add residual blocks at each hidden layer
103        for i in range(len(hidden_sizes) - 1):
104            # Add residual blocks
105            for _ in range(num_residual_blocks):
106                layers.append(ResidualBlock(hidden_sizes[i], dropout_rate))
107            
108            # Project to next hidden size
109            layers.append(nn.Linear(hidden_sizes[i], hidden_sizes[i + 1]))
110            layers.append(nn.LayerNorm(hidden_sizes[i + 1]))
111            layers.append(nn.ReLU())
112            layers.append(nn.Dropout(dropout_rate))
113        
114        # Add final residual blocks
115        for _ in range(num_residual_blocks):
116            layers.append(ResidualBlock(hidden_sizes[-1], dropout_rate))
117        
118        # Output projection
119        layers.append(nn.Linear(hidden_sizes[-1], num_classes))
120        
121        self.layers = layers
122        self._initialize_weights()
123    
124    def _initialize_weights(self):
125        for module in self.modules():
126            if isinstance(module, nn.Linear):
127                nn.init.xavier_uniform_(module.weight)
128                if module.bias is not None:
129                    nn.init.constant_(module.bias, 0)
130            elif isinstance(module, nn.BatchNorm1d):
131                nn.init.constant_(module.weight, 1)
132                nn.init.constant_(module.bias, 0)
133    
134    def forward(self, x: torch.Tensor) -> torch.Tensor:
135        for layer in self.layers:
136            x = layer(x)
137        return x

Neural Network for Hepatitis C classification with residual connections.

Parameters
  • input_size (int): Number of input features.
  • hidden_sizes (list of int): List of hidden layer sizes.
  • num_classes (int): Number of output classes.
  • dropout_rate (float): Dropout rate for regularization.
  • num_residual_blocks (int): Number of residual blocks to use.
Attributes
  • layers (nn.ModuleList): List of network layers including residual blocks.
  • input_size (int): Number of input features.
  • num_classes (int): Number of output classes.
HepatitisNet( input_size: int = 12, hidden_sizes: list = [128, 64, 32], num_classes: int = 2, dropout_rate: float = 0.3, num_residual_blocks: int = 2) 
 86    def __init__(self, input_size: int = 12, hidden_sizes: list = [128, 64, 32], 
 87                 num_classes: int = 2, dropout_rate: float = 0.3, num_residual_blocks: int = 2):
 88        super(HepatitisNet, self).__init__()
 89        
 90        self.input_size = input_size
 91        self.num_classes = num_classes
 92
 93        # Build network architecture
 94        layers = nn.ModuleList()
 95        
 96        # Input projection
 97        layers.append(nn.Linear(input_size, hidden_sizes[0]))
 98        layers.append(nn.LayerNorm(hidden_sizes[0]))
 99        layers.append(nn.ReLU())
100        layers.append(nn.Dropout(dropout_rate))
101        
102        # Add residual blocks at each hidden layer
103        for i in range(len(hidden_sizes) - 1):
104            # Add residual blocks
105            for _ in range(num_residual_blocks):
106                layers.append(ResidualBlock(hidden_sizes[i], dropout_rate))
107            
108            # Project to next hidden size
109            layers.append(nn.Linear(hidden_sizes[i], hidden_sizes[i + 1]))
110            layers.append(nn.LayerNorm(hidden_sizes[i + 1]))
111            layers.append(nn.ReLU())
112            layers.append(nn.Dropout(dropout_rate))
113        
114        # Add final residual blocks
115        for _ in range(num_residual_blocks):
116            layers.append(ResidualBlock(hidden_sizes[-1], dropout_rate))
117        
118        # Output projection
119        layers.append(nn.Linear(hidden_sizes[-1], num_classes))
120        
121        self.layers = layers
122        self._initialize_weights()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

input_size
num_classes
layers
def forward(self, x: torch.Tensor) -> torch.Tensor: 
134    def forward(self, x: torch.Tensor) -> torch.Tensor:
135        for layer in self.layers:
136            x = layer(x)
137        return x

Define the computation performed at every call.

Should be overridden by all subclasses.

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def evaluate_model( model: torch.nn.modules.module.Module, test_loader: torch.utils.data.dataloader.DataLoader, device: str = 'cpu') -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: 
140def evaluate_model(model: nn.Module, test_loader: DataLoader, device: str = 'cpu') -> tuple[np.ndarray, np.ndarray, np.ndarray]:
141    """
142    Evaluate the model on the test dataset.
143    
144    Parameters
145    -----------
146    model : nn.Module
147        The trained model to be evaluated.
148    test_loader : DataLoader
149        DataLoader for the test dataset.
150    device : str
151        Device to run the evaluation on (default: 'cpu').
152
153    Returns
154    -----------
155    tuple[np.ndarray, np.ndarray, np.ndarray]
156        - y_true: Ground truth labels.
157        - y_pred: Predicted labels.
158        - y_probs: Predicted probabilities.
159
160    Examples
161    ---------
162    >>> y_true, y_pred, y_probs = evaluate_model(model, test_loader, device='cuda')
163    """
164    model.eval()
165    y_true = []
166    y_pred = []
167    y_probs = []
168    
169    with torch.no_grad():
170        for data, target in test_loader:
171            data, target = data.to(device), target.to(device)
172            output = model(data)
173            probs = torch.softmax(output, dim=1)
174            pred = output.argmax(dim=1)
175            
176            y_true.extend(target.cpu().numpy())
177            y_pred.extend(pred.cpu().numpy())
178            y_probs.extend(probs.cpu().numpy())
179    
180    return np.array(y_true), np.array(y_pred), np.array(y_probs)

Evaluate the model on the test dataset.

Parameters
  • model (nn.Module): The trained model to be evaluated.
  • test_loader (DataLoader): DataLoader for the test dataset.
  • device (str): Device to run the evaluation on (default: 'cpu').
Returns
  • tuple[np.ndarray, np.ndarray, np.ndarray]: - y_true: Ground truth labels.
    • y_pred: Predicted labels.
    • y_probs: Predicted probabilities.
Examples
>>> y_true, y_pred, y_probs = evaluate_model(model, test_loader, device='cuda')
def save_model( model: torch.nn.modules.module.Module, filepath: str, additional_info: dict = None, demo: bool = False) -> None: 
182def save_model(model: nn.Module, filepath: str, additional_info: dict = None, demo: bool = False) -> None:
183    """
184    Save the model to a file.
185
186    Parameters
187    -----------
188    model : nn.Module
189        The model to be saved.
190    filepath : str
191        Path to the file where the model will be saved.
192    additional_info : dict, optional
193        Any additional information to save with the model (e.g., training parameters).
194    demo : bool, optional
195        Whether the model is being saved in a temp location demo mode (default: False).
196
197    Returns
198    -----------
199    str
200        The path to the saved model file.
201
202    Examples
203    ---------
204    >>> save_model(model, 'models/hepatitis_model.pth', {'input_size': 12, 'num_classes': 2})
205    """
206
207    if demo:
208        filepath = os.path.join(tempfile.gettempdir(), 'hepatitis_model.pth')
209    torch.save({
210        'model_state_dict': model.state_dict(),
211        'model_class': model.__class__.__name__,
212        'additional_info': additional_info
213    }, filepath, _use_new_zipfile_serialization=False)
214    print(f"Model saved to: {filepath}")
215    return filepath

Save the model to a file.

Parameters
  • model (nn.Module): The model to be saved.
  • filepath (str): Path to the file where the model will be saved.
  • additional_info (dict, optional): Any additional information to save with the model (e.g., training parameters).
  • demo (bool, optional): Whether the model is being saved in a temp location demo mode (default: False).
Returns
  • str: The path to the saved model file.
Examples
>>> save_model(model, 'models/hepatitis_model.pth', {'input_size': 12, 'num_classes': 2})
def load_model( filepath: str, model_class: type[HepatitisNet] = <class 'HepatitisNet'>, input_size: int = 12) -> tuple[torch.nn.modules.module.Module, dict]: 
217def load_model(filepath: str, model_class: type[HepatitisNet] = HepatitisNet, input_size: int = 12) -> tuple[nn.Module, dict]:
218    """
219    Load a model from a file.
220    
221    Parameters
222    -----------
223    filepath : str
224        Path to the file from which the model will be loaded.
225    model_class : type
226        The class of the model to be loaded (default: HepatitisNet).
227    input_size : int
228        Number of input features (default: 12).
229
230    Returns
231    -----------
232    tuple[nn.Module, dict]
233        - model: The loaded model.
234        - additional_info: Any additional information saved with the model.
235    
236    Examples
237    ---------
238    >>> model, info = load_model('models/hepatitis_model.pth')
239    >>> print(info)
240    """
241    checkpoint = torch.load(filepath, map_location='cpu', weights_only=False)
242    
243    model = model_class(input_size=input_size)
244    model.load_state_dict(checkpoint['model_state_dict'])
245    model.eval()
246    
247    return model, checkpoint.get('additional_info', None)

Load a model from a file.

Parameters
  • filepath (str): Path to the file from which the model will be loaded.
  • model_class (type): The class of the model to be loaded (default: HepatitisNet).
  • input_size (int): Number of input features (default: 12).
Returns
  • tuple[nn.Module, dict]: - model: The loaded model.
    • additional_info: Any additional information saved with the model.
Examples
>>> model, info = load_model('models/hepatitis_model.pth')
>>> print(info)
class TorchWrapper(sklearn.base.BaseEstimator, sklearn.base.ClassifierMixin): 
249class TorchWrapper(BaseEstimator, ClassifierMixin):
250    """
251    A wrapper to make PyTorch models compatible with scikit-learn.
252    
253    Parameters
254    -----------
255    model : HepatitisNet
256        The PyTorch model instance.
257    device : str
258        Device to run inference on ('cpu' or 'cuda').
259    classes : array-like
260        Class labels for the classifier.
261
262    Attributes
263    -----------
264    model : HepatitisNet
265        The PyTorch model instance.
266    device : str
267        Device to run inference on ('cpu' or 'cuda').
268    classes_ : array-like
269        Class labels for the classifier.
270
271    Examples
272    ---------
273    >>> model, _ = load_model('models/hepatitis_model.pth')
274    >>> wrapper = TorchWrapper(model, device='cuda', classes=[0, 1])
275    >>> CalibrationDisplay.from_estimator(wrapper, X_test, y_test, n_bins=10)
276    """
277
278    def __init__(self, model: type[HepatitisNet], device: str, classes: any):
279        self.model = model
280        self.device = device
281        self.classes_ = classes
282    def fit(self, X: np.ndarray, y: np.ndarray) -> TorchWrapper:
283        return self
284    
285    def predict_proba(self, X: np.ndarray) -> np.ndarray:
286        X_tensor = torch.FloatTensor(X).to(self.device)
287        with torch.no_grad():
288            outputs = self.model(X_tensor)
289            probs = torch.softmax(outputs, dim=1).cpu().numpy()
290        return probs
291    
292    def predict(self, X: np.ndarray) -> np.ndarray:
293        return np.argmax(self.predict_proba(X), axis=1)

A wrapper to make PyTorch models compatible with scikit-learn.

Parameters
  • model (HepatitisNet): The PyTorch model instance.
  • device (str): Device to run inference on ('cpu' or 'cuda').
  • classes (array-like): Class labels for the classifier.
Attributes
  • model (HepatitisNet): The PyTorch model instance.
  • device (str): Device to run inference on ('cpu' or 'cuda').
  • classes_ (array-like): Class labels for the classifier.
Examples
>>> model, _ = load_model('models/hepatitis_model.pth')
>>> wrapper = TorchWrapper(model, device='cuda', classes=[0, 1])
>>> CalibrationDisplay.from_estimator(wrapper, X_test, y_test, n_bins=10)
TorchWrapper( model: type[HepatitisNet], device: str, classes: <built-in function any>) 
278    def __init__(self, model: type[HepatitisNet], device: str, classes: any):
279        self.model = model
280        self.device = device
281        self.classes_ = classes
model
device
classes_
def fit(self, X: numpy.ndarray, y: numpy.ndarray) -> TorchWrapper: 
282    def fit(self, X: np.ndarray, y: np.ndarray) -> TorchWrapper:
283        return self
def predict_proba(self, X: numpy.ndarray) -> numpy.ndarray: 
285    def predict_proba(self, X: np.ndarray) -> np.ndarray:
286        X_tensor = torch.FloatTensor(X).to(self.device)
287        with torch.no_grad():
288            outputs = self.model(X_tensor)
289            probs = torch.softmax(outputs, dim=1).cpu().numpy()
290        return probs
def predict(self, X: numpy.ndarray) -> numpy.ndarray: 
292    def predict(self, X: np.ndarray) -> np.ndarray:
293        return np.argmax(self.predict_proba(X), axis=1)
def set_score_request(unknown):

Descriptor for defining set_{method}_request methods in estimators.

New in version 1.3.

Parameters
  • name (str): The name of the method for which the request function should be created, e.g. "fit" would create a set_fit_request function.
  • keys (list of str): A list of strings which are accepted parameters by the created function, e.g. ["sample_weight"] if the corresponding method accepts it as a metadata.
  • validate_keys (bool, default=True): Whether to check if the requested parameters fit the actual parameters of the method.
Notes

This class is a descriptor 1 and uses PEP-362 to set the signature of the returned function 2.

References

  1. https://docs.python.org/3/howto/descriptor.html ↩

  2. https://www.python.org/dev/peps/pep-0362/ ↩