gaitsetpy.classification.models

View Source

 1from .random_forest import RandomForestModel
 2from .mlp import MLPModel
 3
 4# Optional PyTorch-dependent imports
 5try:
 6    from .lstm import LSTMModel
 7    from .bilstm import BiLSTMModel
 8    from .gnn import GNNModel
 9    from .cnn import CNNModel
10    PYTORCH_AVAILABLE = True
11except ImportError:
12    LSTMModel = None
13    BiLSTMModel = None
14    GNNModel = None
15    CNNModel = None
16    PYTORCH_AVAILABLE = False
17
18def get_classification_model(name: str, **kwargs):
19    """
20    Factory function to get a classification model by name.
21
22    Args:
23        name (str): Name of the model. One of: 'random_forest', 'mlp', 'lstm', 'bilstm', 'gnn', 'cnn'.
24        **kwargs: Model-specific parameters.
25
26    Returns:
27        An instance of the requested model.
28
29    Raises:
30        ValueError: If the model name is not recognized or PyTorch is not available for PyTorch models.
31
32    Example:
33        model = get_classification_model('cnn', input_channels=20, num_classes=4)
34    """
35    name = name.lower()
36    if name == 'random_forest':
37        return RandomForestModel(**kwargs)
38    elif name == 'mlp':
39        return MLPModel(**kwargs)
40    elif name == 'lstm':
41        if not PYTORCH_AVAILABLE or LSTMModel is None:
42            raise ImportError("LSTM model requires PyTorch. Please install PyTorch to use this model.")
43        return LSTMModel(**kwargs)
44    elif name == 'bilstm':
45        if not PYTORCH_AVAILABLE or BiLSTMModel is None:
46            raise ImportError("BiLSTM model requires PyTorch. Please install PyTorch to use this model.")
47        return BiLSTMModel(**kwargs)
48    elif name == 'gnn':
49        if not PYTORCH_AVAILABLE or GNNModel is None:
50            raise ImportError("GNN model requires PyTorch. Please install PyTorch to use this model.")
51        return GNNModel(**kwargs)
52    elif name == 'cnn':
53        if not PYTORCH_AVAILABLE or CNNModel is None:
54            raise ImportError("CNN model requires PyTorch. Please install PyTorch to use this model.")
55        return CNNModel(**kwargs)
56    else:
57        available_models = ['random_forest', 'mlp']
58        if PYTORCH_AVAILABLE:
59            available_models.extend(['lstm', 'bilstm', 'gnn', 'cnn'])
60        raise ValueError(f"Unknown model name: {name}. Supported: {available_models}.")

def get_classification_model(name: str, **kwargs): View Source

19def get_classification_model(name: str, **kwargs):
20    """
21    Factory function to get a classification model by name.
22
23    Args:
24        name (str): Name of the model. One of: 'random_forest', 'mlp', 'lstm', 'bilstm', 'gnn', 'cnn'.
25        **kwargs: Model-specific parameters.
26
27    Returns:
28        An instance of the requested model.
29
30    Raises:
31        ValueError: If the model name is not recognized or PyTorch is not available for PyTorch models.
32
33    Example:
34        model = get_classification_model('cnn', input_channels=20, num_classes=4)
35    """
36    name = name.lower()
37    if name == 'random_forest':
38        return RandomForestModel(**kwargs)
39    elif name == 'mlp':
40        return MLPModel(**kwargs)
41    elif name == 'lstm':
42        if not PYTORCH_AVAILABLE or LSTMModel is None:
43            raise ImportError("LSTM model requires PyTorch. Please install PyTorch to use this model.")
44        return LSTMModel(**kwargs)
45    elif name == 'bilstm':
46        if not PYTORCH_AVAILABLE or BiLSTMModel is None:
47            raise ImportError("BiLSTM model requires PyTorch. Please install PyTorch to use this model.")
48        return BiLSTMModel(**kwargs)
49    elif name == 'gnn':
50        if not PYTORCH_AVAILABLE or GNNModel is None:
51            raise ImportError("GNN model requires PyTorch. Please install PyTorch to use this model.")
52        return GNNModel(**kwargs)
53    elif name == 'cnn':
54        if not PYTORCH_AVAILABLE or CNNModel is None:
55            raise ImportError("CNN model requires PyTorch. Please install PyTorch to use this model.")
56        return CNNModel(**kwargs)
57    else:
58        available_models = ['random_forest', 'mlp']
59        if PYTORCH_AVAILABLE:
60            available_models.extend(['lstm', 'bilstm', 'gnn', 'cnn'])
61        raise ValueError(f"Unknown model name: {name}. Supported: {available_models}.")

Factory function to get a classification model by name.

Args: name (str): Name of the model. One of: 'random_forest', 'mlp', 'lstm', 'bilstm', 'gnn', 'cnn'. **kwargs: Model-specific parameters.

Returns: An instance of the requested model.

Raises: ValueError: If the model name is not recognized or PyTorch is not available for PyTorch models.

Example: model = get_classification_model('cnn', input_channels=20, num_classes=4)