Skip to content

base_trainer

BaseTrainer

Base trainer class for model training.

Parameters:

Name Type Description Default
X np.ndarray or torch.Tensor

Input data of shape (num_samples, num_features).

 required
y np.ndarray or torch.Tensor

Target data of shape (num_samples,).

 required
sample_weights np.ndarray or torch.Tensor

Optional sample weights of shape (num_samples,).

 None
test_size float

The proportion of the data to include in the validation set.

 0.2
random_state int

The seed used by the random number generator.

 42
num_epochs int

Number of training epochs.

 50
batch_size int

Batch size for training.

 256
optimizer_fn_name str

Name of the optimizer function from the torch.optim module.

 'Adam'
lr float

Learning rate for the optimizer.

 0.01
use_scheduler bool

Whether to use a learning rate scheduler (not yet fully supported).

 False
patience int

Number of consecutive epochs with no improvement after which training will be stopped.

 10
dtype torch.dtype

Data type to use for the tensors.

 torch.float32
device torch.device

Device can be specified to the desired cpu or cuda for GPU (else if set to None, then GPU if available, otherwise CPU).

 None
multi_cpu_dataloader bool

Use multiple CPUs for data loading or pass everything onto the GPU.

 True

Attributes:

Name Type Description
X torch.Tensor

Input data tensor.
y torch.Tensor

Target data tensor.
sample_weights torch.Tensor

Sample weights tensor.
test_size float

Proportion of data to include in the validation set.
random_state int

Seed used by the random number generator.
num_epochs int

Number of training epochs.
batch_size int

Batch size for training.
optimizer_fn_name str

Name of the optimizer function.
lr float

Learning rate for the optimizer.
patience int

Number of consecutive epochs with no improvement after which training will be stopped.
dtype torch.dtype

Data type of the tensors.
train_loader torch.utils.data.DataLoader

DataLoader for training data.
Source code in uncertaintyplayground/trainers/base_trainer.py 
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
class BaseTrainer:
    """
    Base trainer class for model training.

    Args:
        X (np.ndarray or torch.Tensor): Input data of shape (num_samples, num_features).
        y (np.ndarray or torch.Tensor): Target data of shape (num_samples,).
        sample_weights (np.ndarray or torch.Tensor): Optional sample weights of shape (num_samples,).
        test_size (float): The proportion of the data to include in the validation set.
        random_state (int): The seed used by the random number generator.
        num_epochs (int): Number of training epochs.
        batch_size (int): Batch size for training.
        optimizer_fn_name (str): Name of the optimizer function from the `torch.optim` module.
        lr (float): Learning rate for the optimizer.
        use_scheduler (bool): Whether to use a learning rate scheduler (`not yet fully supported`).
        patience (int): Number of consecutive epochs with no improvement after which training will be stopped.
        dtype (torch.dtype): Data type to use for the tensors.
        device (torch.device): Device can be specified to the desired `cpu` or `cuda` for GPU (else if set to `None`, then GPU if available, otherwise CPU).
        multi_cpu_dataloader (bool): Use multiple CPUs for data loading or pass everything onto the GPU.

    Attributes:
        X (torch.Tensor): Input data tensor.
        y (torch.Tensor): Target data tensor.
        sample_weights (torch.Tensor): Sample weights tensor.
        test_size (float): Proportion of data to include in the validation set.
        random_state (int): Seed used by the random number generator.
        num_epochs (int): Number of training epochs.
        batch_size (int): Batch size for training.
        optimizer_fn_name (str): Name of the optimizer function.
        lr (float): Learning rate for the optimizer.
        patience (int): Number of consecutive epochs with no improvement after which training will be stopped.
        dtype (torch.dtype): Data type of the tensors.
        train_loader (torch.utils.data.DataLoader): DataLoader for training data.
    """

    def __init__(
            self,
            X,
            y,
            sample_weights=None,
            test_size=0.2,
            random_state=42,
            num_epochs=50,
            batch_size=256,
            optimizer_fn_name="Adam",
            lr=0.01,
            use_scheduler=False,
            patience=10,
            dtype=torch.float32,
            device = None,
            multi_cpu_dataloader = True
    ):
        self.X = X
        self.y = y
        self.sample_weights = sample_weights
        self.test_size = test_size
        self.random_state = random_state
        self.num_epochs = num_epochs
        self.batch_size = batch_size
        self.optimizer_fn_name = optimizer_fn_name
        self.lr = lr
        self.patience = patience
        self.dtype = dtype
        self.use_scheduler = use_scheduler

        # Setting for early stopping
        self.best_epoch = -1
        self.best_val_mse = np.inf

        # Choose device (GPU if available, otherwise CPU)
        self.device = torch.device(
            device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
        )

        # Choose whether to use multiple CPUs for data loading
        self.multi_cpu_dataloader = multi_cpu_dataloader

        # Convert input tensors to the correct type
        self.prepare_inputs()

        # Split data into training and validation sets
        self.split_data()

        # Create DataLoader for training data
        self.prepare_dataloader()

    def prepare_inputs(self):
        """
        Convert input data to the correct type and format.
        """
        # Convert X to a tensor if it's a numpy array
        if isinstance(self.X, np.ndarray):
            self.X = torch.from_numpy(self.X).to(dtype=self.dtype)

        # Convert y to a tensor if it's a list or numpy array
        if isinstance(self.y, (list, np.ndarray)):
            self.y = torch.tensor(self.y, dtype=self.dtype)

        # Check if sample_weights is a tensor, numpy array, or list
        if self.sample_weights is not None:
            if isinstance(self.sample_weights, (np.ndarray, list)):
                self.sample_weights = torch.tensor(
                    self.sample_weights, dtype=self.dtype)

    def split_data(self):
        """
        Split the data into training and validation sets.

        Args:
            test_size (float): Proportion of data to include in the validation set.
        """
        if self.sample_weights is None:
            self.sample_weights = torch.ones(self.X.shape[0], dtype=self.dtype)

        self.X_train, self.X_val, self.y_train, self.y_val, self.sample_weights_train, self.sample_weights_val = \
            train_test_split(self.X, self.y, self.sample_weights,
                             test_size=self.test_size, random_state=self.random_state)

    def custom_lr_scheduler(self, epoch):
        """
        Custom learning rate scheduler function.

        Args:
            epoch (int): Current epoch.

        Returns:
            float: Learning rate for the epoch.
        """
        if epoch < 3:
            return 1 - 0.1 * epoch / self.lr
        else:
            return 0.2 / self.lr

    def prepare_dataloader(self):
            """
            Prepare the DataLoader for training data.
            """
            # Use all available CPU cores or default to 1 if not detected
            if self.multi_cpu_dataloader:
                num_workers = os.cpu_count() - 1 or 1
            else:
                num_workers = 0
            train_dataset = TensorDataset(
                self.X_train, self.y_train, self.sample_weights_train)
            self.train_loader = DataLoader(
                train_dataset, batch_size=self.batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)

custom_lr_scheduler(epoch)

Custom learning rate scheduler function.

Parameters:

Name Type Description Default
epoch int

Current epoch.

 required

Returns:

Name Type Description
float

Learning rate for the epoch.
Source code in uncertaintyplayground/trainers/base_trainer.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def custom_lr_scheduler(self, epoch):
    """
    Custom learning rate scheduler function.

    Args:
        epoch (int): Current epoch.

    Returns:
        float: Learning rate for the epoch.
    """
    if epoch < 3:
        return 1 - 0.1 * epoch / self.lr
    else:
        return 0.2 / self.lr

prepare_dataloader()

Prepare the DataLoader for training data.

Source code in uncertaintyplayground/trainers/base_trainer.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
def prepare_dataloader(self):
        """
        Prepare the DataLoader for training data.
        """
        # Use all available CPU cores or default to 1 if not detected
        if self.multi_cpu_dataloader:
            num_workers = os.cpu_count() - 1 or 1
        else:
            num_workers = 0
        train_dataset = TensorDataset(
            self.X_train, self.y_train, self.sample_weights_train)
        self.train_loader = DataLoader(
            train_dataset, batch_size=self.batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)

prepare_inputs()

Convert input data to the correct type and format.

Source code in uncertaintyplayground/trainers/base_trainer.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def prepare_inputs(self):
    """
    Convert input data to the correct type and format.
    """
    # Convert X to a tensor if it's a numpy array
    if isinstance(self.X, np.ndarray):
        self.X = torch.from_numpy(self.X).to(dtype=self.dtype)

    # Convert y to a tensor if it's a list or numpy array
    if isinstance(self.y, (list, np.ndarray)):
        self.y = torch.tensor(self.y, dtype=self.dtype)

    # Check if sample_weights is a tensor, numpy array, or list
    if self.sample_weights is not None:
        if isinstance(self.sample_weights, (np.ndarray, list)):
            self.sample_weights = torch.tensor(
                self.sample_weights, dtype=self.dtype)

split_data()

Split the data into training and validation sets.

Parameters:

Name Type Description Default
test_size float

Proportion of data to include in the validation set.

 required
Source code in uncertaintyplayground/trainers/base_trainer.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
def split_data(self):
    """
    Split the data into training and validation sets.

    Args:
        test_size (float): Proportion of data to include in the validation set.
    """
    if self.sample_weights is None:
        self.sample_weights = torch.ones(self.X.shape[0], dtype=self.dtype)

    self.X_train, self.X_val, self.y_train, self.y_val, self.sample_weights_train, self.sample_weights_val = \
        train_test_split(self.X, self.y, self.sample_weights,
                         test_size=self.test_size, random_state=self.random_state)