r""" The model consists of a sparse part and a dense part. 1) The dense part is an nn.Linear module that is replicated across all trainers using DistributedDataParallel. 2) The sparse part is a Remote Module that holds an nn.EmbeddingBag on the parameter server. This remote model can
| 17 | |
| 18 | # BEGIN hybrid_model |
| 19 | class HybridModel(torch.nn.Module): |
| 20 | r""" |
| 21 | The model consists of a sparse part and a dense part. |
| 22 | 1) The dense part is an nn.Linear module that is replicated across all trainers using DistributedDataParallel. |
| 23 | 2) The sparse part is a Remote Module that holds an nn.EmbeddingBag on the parameter server. |
| 24 | This remote model can get a Remote Reference to the embedding table on the parameter server. |
| 25 | """ |
| 26 | |
| 27 | def __init__(self, remote_emb_module, device): |
| 28 | super(HybridModel, self).__init__() |
| 29 | self.remote_emb_module = remote_emb_module |
| 30 | self.fc = DDP(torch.nn.Linear(16, 8).cuda(device), device_ids=[device]) |
| 31 | self.device = device |
| 32 | |
| 33 | def forward(self, indices, offsets): |
| 34 | emb_lookup = self.remote_emb_module.forward(indices, offsets) |
| 35 | return self.fc(emb_lookup.cuda(self.device)) |
| 36 | # END hybrid_model |
| 37 | |
| 38 | # BEGIN setup_trainer |