Sorry actually I just realized I’ve completely misspoken. I’m not wrapping my model in DDP. I was planning on doing this and then I realized it replicates across nodes where as I need to send distinct (but related) models to nodes.



 mrshenli:

Looks like you are doing inference instead of training? In this case, don’t you need to somehow gather/combine the outputs from the four different Conv2d layers from 4 different DDP processes? Otherwise, how did you get the final inference result?


Yes this is correct, I do a map and then I plan on doing a concat to reconstruct the output channels as if they all came from the same Conv2d. I think you can basically get the idea from this code snippet
def run(rank, size):
    with torch.no_grad():
        image_pth = Path(os.path.dirname(os.path.realpath(__file__))) / Path(
            "../simulation/screenshot.png"
        )

        screenshot = SimulPLIF(img_path=image_pth, num_repeats=1, load_truth=False)
        img_height, img_width = screenshot[0].squeeze(0).numpy().shape

        from nn_dog import PIN_MEMORY

        train_dataloader = DataLoader(screenshot, batch_size=1, pin_memory=PIN_MEMORY)
        dog = DifferenceOfGaussiansFFT(
            img_height=img_height,
            img_width=img_width,
            sigma_bins=48 // size,
            max_sigma=30,
        ).to(rank, non_blocking=PIN_MEMORY)
        for p in dog.parameters():
            p.requires_grad = False
        dog.eval()
        torch.cuda.synchronize(rank)

        dogs = []
        for i in range(10):
            img_tensor = next(iter(train_dataloader))
            img_tensor = img_tensor.to(rank)
            torch.cuda.synchronize(rank)
            dogs.append(dog(img_tensor))
        return dogs


def init_process(rank_size_fn, backend="nccl"):
    rank, size, fn = rank_size_fn
    """ Initialize the distributed environment. """
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29500"
    dist.init_process_group(backend, rank=rank, world_size=size)
    return fn(rank, size)


if __name__ == "__main__":
    set_start_method("spawn")

    size = 4
    pool = Pool(processes=size)
    start = time.monotonic()
    res = pool.map(init_process, [(i, size, run) for i in range(size)])
    end = time.monotonic()
    print(end - start)