Train the model Parameters ---------- epochs: int number of epochs to train lr: float learning rate weight_decay: float weight decay Returns ------- None
(self, epochs = 20,
lr = 1e-3,
weight_decay = 5e-4
)
| 395 | |
| 396 | |
| 397 | def train(self, epochs = 20, |
| 398 | lr = 1e-3, |
| 399 | weight_decay = 5e-4 |
| 400 | ): |
| 401 | """ |
| 402 | Train the model |
| 403 | |
| 404 | Parameters |
| 405 | ---------- |
| 406 | epochs: int |
| 407 | number of epochs to train |
| 408 | lr: float |
| 409 | learning rate |
| 410 | weight_decay: float |
| 411 | weight decay |
| 412 | |
| 413 | Returns |
| 414 | ------- |
| 415 | None |
| 416 | |
| 417 | """ |
| 418 | |
| 419 | train_loader = self.dataloader['train_loader'] |
| 420 | val_loader = self.dataloader['val_loader'] |
| 421 | |
| 422 | self.model = self.model.to(self.device) |
| 423 | best_model = deepcopy(self.model) |
| 424 | optimizer = optim.Adam(self.model.parameters(), lr=lr, weight_decay = weight_decay) |
| 425 | scheduler = StepLR(optimizer, step_size=1, gamma=0.5) |
| 426 | |
| 427 | min_val = np.inf |
| 428 | print_sys('Start Training...') |
| 429 | |
| 430 | for epoch in range(epochs): |
| 431 | self.model.train() |
| 432 | |
| 433 | for step, batch in enumerate(train_loader): |
| 434 | batch.to(self.device) |
| 435 | optimizer.zero_grad() |
| 436 | y = batch.y |
| 437 | if self.config['uncertainty']: |
| 438 | pred, logvar = self.model(batch) |
| 439 | loss = uncertainty_loss_fct(pred, logvar, y, batch.pert, |
| 440 | reg = self.config['uncertainty_reg'], |
| 441 | ctrl = self.ctrl_expression, |
| 442 | dict_filter = self.dict_filter, |
| 443 | direction_lambda = self.config['direction_lambda']) |
| 444 | else: |
| 445 | pred = self.model(batch) |
| 446 | loss = loss_fct(pred, y, batch.pert, |
| 447 | ctrl = self.ctrl_expression, |
| 448 | dict_filter = self.dict_filter, |
| 449 | direction_lambda = self.config['direction_lambda']) |
| 450 | loss.backward() |
| 451 | nn.utils.clip_grad_value_(self.model.parameters(), clip_value=1.0) |
| 452 | optimizer.step() |
| 453 | |
| 454 | if self.wandb: |