I have a Siamese Network with a triplet loss function at the end. E.g.:
> class Siamese(nn.Module): > def __init__(self, ae_net): > super(Siamese, self).__init__() > self.ae_net = ae_net > > def forward(self, x1, x2, x3, hidden): > > a = self.ae_net(x1, hidden) > b = self.ae_net(x2, hidden) > c = self.ae_net(x3, hidden) > return a, b, c
The network that repeats itself (i.e.,
self.ae_net) is an LSTM with inputs of varying lengths, so I’m not sure I can use
nn.DataParallel. I was wondering if there was a way to assign the implementation of every instance of
self.ae_net to a different GPU, so that they would be calculated in parallel.