(self, task, i_task, epochs)
| 100 | return loss |
| 101 | |
| 102 | def train_one_task(self, task, i_task, epochs): |
| 103 | print('task = ', task) |
| 104 | |
| 105 | dataloader_train = self.train_task_list[task] |
| 106 | self.train_length = len(dataloader_train) |
| 107 | total_steps = epochs |
| 108 | progress_bar = tqdm(total=total_steps, leave=True, disable=(self.args.global_rank != 0)) |
| 109 | |
| 110 | for epoch in range(epochs): |
| 111 | print(epoch) |
| 112 | self.model.train() |
| 113 | |
| 114 | for step, batch in enumerate(tqdm(dataloader_train)): |
| 115 | del batch['sources'] |
| 116 | batch = {k: batch[k].to('cuda') for k in batch} |
| 117 | loss = self.train_step(batch, task) |
| 118 | |
| 119 | if self.args.global_rank == 0: |
| 120 | progress_bar.update(1) |
| 121 | description = f"Epoch {epoch + 1}, Step {step}, Loss: {loss.item():.4f}" |
| 122 | progress_bar.set_description(description, refresh=False) |
| 123 | |
| 124 | self.model.backward(loss) |
| 125 | self.model.step() |
| 126 | |
| 127 | def evaluate_one_task(self, round, infer_task_id, task): |
| 128 | if self.args.local_rank == -1: |
nothing calls this directly
no test coverage detected