I’m working on retraining CLIP to train on several new tasks.
I’m still new to dataparallel(I’m not using the DDP since DP seems simpler to implement).
This question is related to this thread
  
    
    
    How to create a network with multiple forward routines 
  
  
    Hi, 
I would do the following: 
class H:
  def __init__(self, *args):
    # init H

  def forward(self, inp):
    # forward of H
    return out

class G:
  def __init__(self, *args):
    # init G

  def forward(self, inp):
    # forward of G
    return out

class F:
  def __init__(self, *args):
    # init F
    self.g = G()
    self.h = H()

  def forward(self, inp):
    # forward of F
    return self.g(self.h(inp))

model = F()

# forward the whole model
model(inp)

# foward only g
model.g(inp)…
  


Since I want to use data parallel on non-forward method of the network, I decide to write this :
class CustomModel1(nn.Module):
  def __init__(self, model) :
        super(CustomModel1, self).__init__()
        self.model = model
  def forward(self,x,y):
    output1 = self.model.custom_methodA(x)
    output2 = self.model.custom_methodA(y)
    <do_something>
    return result1,result2

class CustomModel2(nn.Module):
  def __init__(self, model) :
        super(CustomModel2, self).__init__()
        self.model = model
  def forward(self,x,y):
    output1 = self.model.custom_methodB(x)
    output2 = self.model.custom_methodB(y)
    <do_something>
    return result1,result2

model = CLIP()
model1 = CustomModel1(model)
model2 = CustomModel2(model)

model1 = nn.DataParallel(model1)
model2 = nn.DataParallel(model2)

optimizer = torch.optim.Adam(model.parameters())
loss1 = nn.BCELoss()
loss2 = nn.BCELoss()

for batch in dataloader:
  a,b,c,d,y = batch
  result1 = model1(a,b)
  result2 = model2(c,d)
  loss = loss1(result1,y) + loss2(result2,y)
  loss.backward()
  optimizer.step()
  optimizer.zero_grad()

This code will give RuntimeError: CUDA error: device-side assert triggered at the loss.backward().
Which step is wrong in my code? I feel some stuff is wrong but can’t figure out how to handle it the right way.  Like, is the computational graph of CLIP inside CustomModel2 and 1 still connected to the original model? Should my optimizer connected to model, or should I create separate optimizer for model1 and model2?
Thanks