Многопроцессорность / DDP - Барьерные блоки loss.backward () - PullRequest
Новая
       95

Многопроцессорность / DDP - Барьерные блоки loss.backward ()

0 голосов
/ 08 мая 2020

У меня есть пример игрушки для DDP. У меня два процесса, и один процесс имеет больше «итераций», чем другой 

    if self.id == 0:
        ep = 10
        sl = 0.1
    elif self.id == 1:
        ep = 15
        sl = 0.2

    for i in range(0, ep):

        # if i == 10:
        #     break

        self._log.info(self.id, "image.. " + str(i))
        image = torch.zeros((1,2,3,256,384))
        input = {"rgb": image}
        self._log.info(self.id, "model.. " + str(i))
        output = self.model(input)
        loss = torch.sum(output["output"][-1])


        self._log.info(self.id, "zerograd.. " + str(i))
        self.optimizer.zero_grad()
        self._log.info(self.id, "backward.. " + str(i))
        loss.backward()
        self._log.info(self.id, "step.. " + str(i))
        self.optimizer.step()
        print("step: " + str(self.id) + " " + str(i))
        self._log.info(self.id, "done " + str(i))

        time.sleep(sl)

    self._log.info(self.id, "barrier.. ")
    dist.barrier()

Но как только первый процесс (с меньшим количеством итераций) достигает dist.barrier (), loss.backward () становится заблокирован во 2-м процессе! Как мне обойти это? Я мог бы заставить все процессы иметь одинаковое количество итераций ... но это не то, чего я хочу.

...