| 29 | |
| 30 | |
| 31 | class BatchUpdateParameterServer(object): |
| 32 | |
| 33 | def __init__(self, batch_update_size=batch_update_size): |
| 34 | self.model = torchvision.models.resnet50(num_classes=num_classes) |
| 35 | self.lock = threading.Lock() |
| 36 | self.future_model = torch.futures.Future() |
| 37 | self.batch_update_size = batch_update_size |
| 38 | self.curr_update_size = 0 |
| 39 | self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9) |
| 40 | for p in self.model.parameters(): |
| 41 | p.grad = torch.zeros_like(p) |
| 42 | |
| 43 | def get_model(self): |
| 44 | return self.model |
| 45 | |
| 46 | @staticmethod |
| 47 | @rpc.functions.async_execution |
| 48 | def update_and_fetch_model(ps_rref, grads): |
| 49 | self = ps_rref.local_value() |
| 50 | timed_log(f"PS got {self.curr_update_size}/{batch_update_size} updates") |
| 51 | for p, g in zip(self.model.parameters(), grads): |
| 52 | p.grad += g |
| 53 | with self.lock: |
| 54 | self.curr_update_size += 1 |
| 55 | fut = self.future_model |
| 56 | |
| 57 | if self.curr_update_size >= self.batch_update_size: |
| 58 | for p in self.model.parameters(): |
| 59 | p.grad /= self.batch_update_size |
| 60 | self.curr_update_size = 0 |
| 61 | self.optimizer.step() |
| 62 | self.optimizer.zero_grad(set_to_none=False) |
| 63 | fut.set_result(self.model) |
| 64 | timed_log("PS updated model") |
| 65 | self.future_model = torch.futures.Future() |
| 66 | |
| 67 | return fut |
| 68 | |
| 69 | |
| 70 | class Trainer(object): |