| 68 | |
| 69 | |
| 70 | class Trainer(object): |
| 71 | |
| 72 | def __init__(self, ps_rref): |
| 73 | self.ps_rref = ps_rref |
| 74 | self.loss_fn = nn.MSELoss() |
| 75 | self.one_hot_indices = torch.LongTensor(batch_size) \ |
| 76 | .random_(0, num_classes) \ |
| 77 | .view(batch_size, 1) |
| 78 | |
| 79 | def get_next_batch(self): |
| 80 | for _ in range(num_batches): |
| 81 | inputs = torch.randn(batch_size, 3, image_w, image_h) |
| 82 | labels = torch.zeros(batch_size, num_classes) \ |
| 83 | .scatter_(1, self.one_hot_indices, 1) |
| 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): |