my_project.train

  1"""Training utilities for the automobile price prediction model."""
  2
  3from pathlib import Path
  4
  5import torch
  6import torch.nn as nn
  7from torch.utils.data import DataLoader
  8
  9from my_project.dataset import AutomobileDataset
 10
 11
 12def create_model(input_size: int) -> nn.Module:
 13
 14    """Create the neural network used for price prediction.
 15
 16    Parameters
 17    ----------
 18    input_size : int
 19        Number of input features.
 20
 21    Returns
 22    -------
 23    torch.nn.Module
 24        Neural network for predicting automobile selling prices.
 25
 26    """
 27
 28    return nn.Sequential(
 29        nn.Linear(input_size, 32),
 30        nn.ReLU(),
 31        nn.Linear(32, 1)
 32    )
 33
 34
 35def train_model(
 36    model: nn.Module,
 37    train_loader: DataLoader,
 38    epochs: int = 100,
 39    learning_rate: float = 0.001
 40) -> nn.Module:
 41
 42    """Train the automobile price prediction model.
 43
 44    Parameters
 45    ----------
 46    model : torch.nn.Module
 47        Neural network to train.
 48
 49    train_loader : torch.utils.data.DataLoader
 50        DataLoader containing the training data.
 51
 52    epochs : int, default=100
 53        Number of training epochs.
 54
 55    learning_rate : float, default=0.001
 56        Learning rate used by the Adam optimizer.
 57
 58    Returns
 59    -------
 60    torch.nn.Module
 61        The trained model.
 62
 63    """
 64
 65    loss_fn = nn.MSELoss()
 66
 67    optimizer = torch.optim.Adam(
 68        model.parameters(),
 69        lr=learning_rate
 70    )
 71
 72    for epoch in range(epochs):
 73
 74        for X, y in train_loader:
 75            prediction = model(X).squeeze()
 76            loss = loss_fn(prediction, y)
 77
 78            optimizer.zero_grad()
 79            loss.backward()
 80            optimizer.step()
 81
 82        print(f"Epoch {epoch + 1}, Loss: {loss.item():.4f}")
 83
 84    return model
 85
 86
 87def main():
 88    """Train the automobile price prediction model and save it to disk."""
 89
 90    PROCESSED_TRAIN_FILE = "data/processed/automobile_dataset"
 91    PROCESSED_TEST_FILE = "data/processed/automobile_test"
 92    MODEL_PATH = Path("models/automobile_model.pth")
 93
 94    MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
 95
 96    train_dataset = AutomobileDataset(
 97        f"{PROCESSED_TRAIN_FILE}.parquet"
 98    )
 99
100    test_dataset = AutomobileDataset(
101        f"{PROCESSED_TEST_FILE}.parquet"
102    )
103
104    train_loader = DataLoader(
105        train_dataset,
106        batch_size=32,
107        shuffle=True
108    )
109
110    # The test loader is prepared here for later evaluation.
111    test_loader = DataLoader(
112        test_dataset,
113        batch_size=32,
114        shuffle=False
115    )
116
117    model = create_model(train_dataset.X.shape[1])
118
119    model = train_model(
120        model,
121        train_loader
122    )
123
124    torch.save(
125        model.state_dict(),
126        MODEL_PATH
127    )
128
129    print(f"Model saved to {MODEL_PATH}")
130
131
132if __name__ == "__main__":
133    main()