Commit dd0ee581 authored by Lorenzo, Leodegario U. II's avatar Lorenzo, Leodegario U. II
Browse files

Merge branch 'class-weights' into 'master'

Add class weights

See merge request !1
parents 9348a9fb d2e83331
......@@ -567,7 +567,7 @@ class MLModels:
def train_test(self, X, y, n_trials=100, scorer=None,
stratify=False, smote=False, test_size=0.25,
scaling=None):
scaling=None, class_weight=None):
"""
Calculate the train and test accuracy given input and target features
......@@ -594,6 +594,9 @@ class MLModels:
scaling : str, default=None
Specify the scaling to be performed. By default does not scale the
data. See MLModels.scalers() for available options.
class_weight : dict or 'balanced', default=None
Weights associated with classes in the form
`{class_label: weight}`
"""
# Initialize results container
train_accuracies = []
......@@ -668,11 +671,24 @@ class MLModels:
# Get start time
start_time = time.time()
# Initialize and train the model
# Initialize model
if self._setting_name is not None:
clf = self.model(**{self._setting_name: s})
else:
clf = self.model
# Set class weights
if class_weight is not None:
# Get class parameters
params = clf.get_params()
# Check if classifier has class weights attribute
if 'class_weight' in params.keys():
# Set class weight
params['class_weight'] = class_weight
clf.set_params(**params)
# Train model
clf.fit(X_train, y_train)
# Get train time
......@@ -736,7 +752,7 @@ class MLModels:
n_neighbors=None, use_methods=None, n_trials=100,
tree_rs=None, scorer=None, task='C',
stratify=False, smote=False, test_size=0.25,
scaling=None):
scaling=None, class_weight=None):
"""
Perform accuracy measurements for various models given dataset
......@@ -781,6 +797,9 @@ class MLModels:
scaling : str, default=None
Specify the scaling to be performed. By default does not scale the
data. See MLModels.scalers() for available options.
class_weight : dict or 'balanced', default=None
Weights associated with classes in the form
`{class_label: weight}`
Returns
-------
......@@ -850,8 +869,9 @@ class MLModels:
# Train method
m.train_test(X, labels, n_trials=n_trials, scorer=scorer,
stratify=stratify, smote=smote,
test_size=test_size, scaling=scaling)
test_size=test_size, scaling=scaling,
class_weight=class_weight)
# Set feature names
if feature_names is not None:
m.feature_names = np.array(feature_names)
......
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment