How to run network with multiple independent inputs in parallel in Pytorch?

distributed
1

I have a neural network with two separate vectors as inputs, similar to this question. Both inputs are encoded and then processed further. But until then, the encoding of those inputs is completely independent. How can I parallelize the encoding phase in pytorch? A minimal example of my code:

class MyModel(nn.Module):
    def __init__(params):
        super(MyModel self).__init__()

        self.encoder1 = Encoder1()
        self.encoder2 = Encoder2()

    def forward(x1, x2):
        # how to calculate both encodings in parallel?
        enc1 = self.encoder1(x1)
        enc2 = self.encoder2(x2)
        return some_func(enc1, enc2)

I checked pytorch forums and learned a bit about DataParallel, but i am not sure how to fit it to my case and on only 1 GPU.

2

DataParallel won’t help here if there is only one GPU, but you could use multiple CUDA streams. For example,

s0 = torch.cuda.Stream()
s1 = torch.cuda.Stream()
with torch.cuda.stream(s0):
    enc1 = self.encoder1(x1)

with torch.cuda.stream(s1):
    enc2 = self.encoder2(x2)
3

Thank you very much! :slight_smile: