Train the model and return training history
(model, train_loader, val_loader, epochs=100, lr=0.001, device='cpu')
| 48 | |
| 49 | |
| 50 | def train_model(model, train_loader, val_loader, epochs=100, lr=0.001, device='cpu'): |
| 51 | """Train the model and return training history""" |
| 52 | model = model.to(device) |
| 53 | criterion = nn.MSELoss() |
| 54 | optimizer = optim.Adam(model.parameters(), lr=lr) |
| 55 | |
| 56 | history = {'train_loss': [], 'val_loss': []} |
| 57 | |
| 58 | for epoch in range(epochs): |
| 59 | # Training |
| 60 | model.train() |
| 61 | train_losses = [] |
| 62 | for X_batch, y_batch in train_loader: |
| 63 | X_batch, y_batch = X_batch.to(device), y_batch.to(device) |
| 64 | optimizer.zero_grad() |
| 65 | pred = model(X_batch) |
| 66 | loss = criterion(pred, y_batch) |
| 67 | loss.backward() |
| 68 | optimizer.step() |
| 69 | train_losses.append(loss.item()) |
| 70 | |
| 71 | # Validation |
| 72 | model.eval() |
| 73 | val_losses = [] |
| 74 | with torch.no_grad(): |
| 75 | for X_batch, y_batch in val_loader: |
| 76 | X_batch, y_batch = X_batch.to(device), y_batch.to(device) |
| 77 | pred = model(X_batch) |
| 78 | loss = criterion(pred, y_batch) |
| 79 | val_losses.append(loss.item()) |
| 80 | |
| 81 | history['train_loss'].append(np.mean(train_losses)) |
| 82 | history['val_loss'].append(np.mean(val_losses)) |
| 83 | |
| 84 | return model, history |
| 85 | |
| 86 | |
| 87 | def evaluate_model(model, X_test, y_test, device='cpu'): |
no test coverage detected