ourflexigrad.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import pdb
import numpy as np
import copy
import random


class Flexigrad():
    def __init__(self, optimizer, reduction='mean'): #optimizer用于优化模型参数的优化器，reduction指定如何合并多个任务的梯度
        self._optim, self._reduction = optimizer, reduction
        return

    #optimizer属性用于返回 PCGrad 类中存储的优化器
    @property
    def optimizer(self):
        return self._optim

    def zero_grad(self):
        '''
        clear the gradient of the parameters
        '''

        return self._optim.zero_grad(set_to_none=True)

    def step(self):
        '''
        update the parameters with the gradient
        '''

        return self._optim.step()

    def pc_backward(self, objectives):
        '''
        calculate the gradient of the parameters

        input:
        - objectives: a list of objectives
        '''

        grads, shapes, has_grads = self._pack_grad(objectives) #获取每个任务的梯度、参数的形状以及是否有梯度的掩码
        
        # for idx, grad in enumerate(grads):
        #     print(f"Task {idx + 1} gradient count: {len(grad)}")
        
        pc_grad = self._project_conflicting(grads, has_grads) #修正冲突的梯度
        pc_grad = self._unflatten_grad(pc_grad, shapes[0]) #将修正后的梯度还原为原始形状
        self._set_grad(pc_grad) #将修正后的梯度设置回模型参数
        return

    def _project_conflicting(self, grads, has_grads, shapes=None):
        shared = torch.stack(has_grads).prod(0).bool() #shared是布尔掩码，计算哪些参数在所有任务中都有梯度
        pc_grad, num_task = copy.deepcopy(grads), len(grads) #复制梯度列表并获取任务数量

        for i in range(num_task):
            for j in range(i + 1, num_task):
                g_i, g_j = pc_grad[i], grads[j]
                
                cos_theta = torch.dot(g_i, g_j) / (g_i.norm() * g_j.norm() + 1e-8)  # 避免除零
                
                if cos_theta < 0:  # 夹角 > 90°，执行投影以减少冲突
                    g_i -= (torch.dot(g_i, g_j) / (g_j.norm() ** 2 + 1e-8)) * g_j
                else:  # 夹角小于90°，计算权重调整
                    w_i = (1 + cos_theta) / 2  # 归一化权重，范围 [0.5, 1]
                    pc_grad[i] = w_i * g_i + (1 - w_i) * ((torch.dot(g_i, g_j) / (g_j.norm() ** 2 + 1e-8)) * g_j)

        merged_grad = torch.zeros_like(grads[0]).to(grads[0].device) #合并修正后的梯度。
        #如果 reduction 为 'mean'，则取平均；如果为 'sum'，则求和
        if self._reduction:
            merged_grad[shared] = torch.stack([g[shared]
                                           for g in pc_grad]).mean(dim=0)
        elif self._reduction == 'sum':
            merged_grad[shared] = torch.stack([g[shared]
                                           for g in pc_grad]).sum(dim=0)
        else: exit('invalid reduction method')

        merged_grad[~shared] = torch.stack([g[~shared]
                                            for g in pc_grad]).sum(dim=0)
        return merged_grad

    #将修正后的梯度设置回模型参数
    def _set_grad(self, grads):
        '''
        set the modified gradients to the network
        '''

        idx = 0 #用于跟踪梯度列表的索引
        for group in self._optim.param_groups: #遍历优化器的参数组和参数，将修正后的梯度赋值给对应的参数
            for p in group['params']:
                # if p.grad is None: continue
                p.grad = grads[idx]
                idx += 1 
        return

    #用于获取每个任务的梯度、参数的形状以及是否有梯度的掩码
    def _pack_grad(self, objectives):
        '''
        pack the gradient of the parameters of the network for each objective
        
        output:
        - grad: a list of the gradient of the parameters
        - shape: a list of the shape of the parameters
        - has_grad: a list of mask represent whether the parameter has gradient
        '''
        grads, shapes, has_grads = [], [], []
        for obj in objectives:
        # for idx, obj in enumerate(objectives):
            self._optim.zero_grad(set_to_none=True)
            obj.backward(retain_graph=True)

            # task_name = ["Order", "Family", "Species"][idx]  # 假设任务名称是 Order, Family, Species
            # print(f"\nGradients for Task {idx + 1} ({task_name}):")
            grad, shape, has_grad = self._retrieve_grad() #获取梯度、形状和是否有梯度的掩码
            grads.append(self._flatten_grad(grad, shape))
            has_grads.append(self._flatten_grad(has_grad, shape))
            shapes.append(shape)
        return grads, shapes, has_grads

    #用于将展平的梯度还原为原始形状
    def _unflatten_grad(self, grads, shapes):
        unflatten_grad, idx = [], 0
        for shape in shapes:
            length = np.prod(shape)
            unflatten_grad.append(grads[idx:idx + length].view(shape).clone())
            idx += length
        return unflatten_grad

    def _flatten_grad(self, grads, shapes):
        flatten_grad = torch.cat([g.flatten() for g in grads])
        return flatten_grad

    #用于获取当前目标函数的梯度、参数的形状以及是否有梯度的掩码
    def _retrieve_grad(self):
        '''
        get the gradient of the parameters of the network with specific 
        objective
        
        output:
        - grad: a list of the gradient of the parameters
        - shape: a list of the shape of the parameters
        - has_grad: a list of mask represent whether the parameter has gradient
        '''

        grad, shape, has_grad = [], [], []
        for group in self._optim.param_groups:
            for p in group['params']:
                # if p.grad is None: continue
                # tackle the multi-head scenario
                
                # print(f"Parameter shape: {p.shape}, Gradient:")
                # print(p.grad)
                if p.grad is None:
                    shape.append(p.shape)
                    grad.append(torch.zeros_like(p).to(p.device))
                    has_grad.append(torch.zeros_like(p).to(p.device))
                    continue
                shape.append(p.grad.shape)
                grad.append(p.grad.clone())
                has_grad.append(torch.ones_like(p).to(p.device))
        return grad, shape, has_grad


class TestNet(nn.Module):
    def __init__(self):
        super().__init__()
        self._linear = nn.Linear(3, 4)

    def forward(self, x):
        return self._linear(x)


class MultiHeadTestNet(nn.Module):
    def __init__(self):
        super().__init__()
        self._linear = nn.Linear(3, 2)
        self._head1 = nn.Linear(2, 4)
        self._head2 = nn.Linear(2, 4)

    def forward(self, x):
        feat = self._linear(x)
        return self._head1(feat), self._head2(feat)


if __name__ == '__main__':

    # fully shared network test
    torch.manual_seed(4)
    x, y = torch.randn(2, 3), torch.randn(2, 4)
    net = TestNet()
    y_pred = net(x)
    pc_adam = Flexigrad(optim.Adam(net.parameters()))
    pc_adam.zero_grad()
    loss1_fn, loss2_fn = nn.L1Loss(), nn.MSELoss()
    loss1, loss2 = loss1_fn(y_pred, y), loss2_fn(y_pred, y)

    pc_adam.pc_backward([loss1, loss2])
    for p in net.parameters():
        print(p.grad)

    print('-' * 80)
    # seperated shared network test

    torch.manual_seed(4)
    x, y = torch.randn(2, 3), torch.randn(2, 4)
    net = MultiHeadTestNet()
    y_pred_1, y_pred_2 = net(x)
    pc_adam = Flexigrad(optim.Adam(net.parameters()))
    pc_adam.zero_grad()
    loss1_fn, loss2_fn = nn.MSELoss(), nn.MSELoss()
    loss1, loss2 = loss1_fn(y_pred_1, y), loss2_fn(y_pred_2, y)

    pc_adam.pc_backward([loss1, loss2])
    for p in net.parameters():
        print(p.grad)