Pytorch MultiTask Model

1

Hello Everyone!
I’m trying to make a MultiTask Regressor Model using pytorch.
However, when the code tries to backward a second time, the error is caused.

I put backward(retain_graph=True) and it works, but i think this works the results through both indivisial layer 1 and 2.
I wanna get result seperately trough shared layer, indivisial layer1 and shared layer,indivisial layer2

Could you guys give some advice to me?

Here is my code.

    loss_ = []  # loss를 저장할 리스트


    for epoch in range(400): 
        running_loss = 0.0  
    
        for _, data in enumerate(trainloader, 0):
            inputs, values = data  
            inputs, values = inputs, values

#  forward                
            shared_result = shared_model(inputs)           
            outputs_1 = indivi_model_1(shared_result)
            outputs_2 = indivi_model_2(shared_result)

# Backward
            # Shared Layer
            # shared_optimizer.zero_grad()        
            # shared_optimizer.step()
            # # loss.backward()                
            
            # Indivisial Layer_1         
            indivisial_optimizer_1.zero_grad()
            loss_1 = criterion(outputs_1, values)  
            loss_1.backward()
            indivisial_optimizer_1.step()
            
            # Indivisial Layer_2
            indivisial_optimizer_2.zero_grad()
            loss_2 = criterion(outputs_2, values)    
            loss_2.backward()
            indivisial_optimizer_2.step()
2

Would this be something you’re looking for? Individually forward on the shared model?

        # Zero all gradients
        shared_optimizer.zero_grad()
        indivisial_optimizer_1.zero_grad()
        indivisial_optimizer_2.zero_grad()

        # Forward pass for task 1
        shared_result_1 = shared_model(inputs)
        outputs_1 = indivi_model_1(shared_result_1)
        loss_1 = criterion(outputs_1, values)

        # Backward pass for task 1
        loss_1.backward(retain_graph=True)

        # Forward pass for task 2
        shared_result_2 = shared_model(inputs)
        outputs_2 = indivi_model_2(shared_result_2)
        loss_2 = criterion(outputs_2, values)

        # Backward pass for task 2
        loss_2.backward()

        # Update weights
        indivisial_optimizer_1.step()
        indivisial_optimizer_2.step()
        shared_optimizer.step()

Here is a colab notebook for a small end-to-end example:

1 Like
3

Thank you so much!
Yes, This is what i’ve been looking for:)