(self)
| 84 | yield inputs.cuda(), labels.cuda() |
| 85 | |
| 86 | def train(self): |
| 87 | name = rpc.get_worker_info().name |
| 88 | m = self.ps_rref.rpc_sync().get_model().cuda() |
| 89 | for inputs, labels in self.get_next_batch(): |
| 90 | timed_log(f"{name} processing one batch") |
| 91 | self.loss_fn(m(inputs), labels).backward() |
| 92 | timed_log(f"{name} reporting grads") |
| 93 | m = rpc.rpc_sync( |
| 94 | self.ps_rref.owner(), |
| 95 | BatchUpdateParameterServer.update_and_fetch_model, |
| 96 | args=(self.ps_rref, [p.grad for p in m.cpu().parameters()]), |
| 97 | ).cuda() |
| 98 | timed_log(f"{name} got updated model") |
| 99 | |
| 100 | |
| 101 | def run_trainer(ps_rref): |
no test coverage detected