This repository provide code of reimplementation for Gradient Surgery for Multi-Task Learning in PyTorch 2.8.0 and a fork from github.com/WeiChengTseng/Pytorch-PCGrad. We modified Wei Tseng's work to make it compatible with PyTorch Distributed Data Parallel (DDP) framework.
The minimum framework requirement for this codebase is Python 3.11. To leverage CUDA, we used the following framework: CUDA 12.1, NVIDIA Driver v 560.35.03 and CuDNN 8.9.2.
We provide an Anaconda environment file (environment.yml) to streamline the setup of the working environment.
conda env create -f environment.yml
conda activate pcgrad
MultiMNIST is a Multi-Task Learning (MTL) adaptation of MNIST.
In order to convert digit classification into a multi-task problem, the original authors overlay multiple images together. We use a similar construction. For each image, a different one is chosen uniformly in random. Then one of these images is put at the top-left and the other one is at the bottom-right. The resulting tasks are: classifying the digit on the top-left (task-L) and classifying the digit on the bottom-right (task-R). We use 60K examples and directly apply existing single-task MNIST models. The original implementation is available at github.com/intel-isl/MultiObjectiveOptimization.
import torch
import torch.nn as nn
import torch.optim as optim
from pcgrad import PCGrad
# wrap your favorite optimizer
optimizer = PCGrad(optim.Adam(model.parameters()))
# a list of per-task losses
losses = [...]
# calculate the gradient can apply gradient modification
optimizer.pc_backward(losses)
# apply gradient step
optimizer.step()As mentioned by Wei Tseng, part of implementation is leveraged from github.com/intel-isl/MultiObjectiveOptimization. Use the following command to run the training pipeline.
python train.py --pcgrad --batch_size 128 --n_epochs 100
| Method | val loss | left-digit | right-digit |
|---|---|---|---|
| Jointly Training | 1.04 | 91.26 | 89.03 |
| PCGrad (sum reduction) | 1.06 | 91.51 | 88.97 |
| PCGrad (mean reduction) | 1.02 | 91.58 | 89.35 |
@inproceedings{yu2020gradient,
title = {Gradient Surgery for Multi-Task Learning},
author = {Tianhe Yu and Saurabh Kumar and Abhishek Gupta and Sergey Levine and Karol Hausman and Chelsea Finn},
editor = {Hugo Larochelle and Marc'Aurelio Ranzato and Raia Hadsell and Maria{-}Florina Balcan and Hsuan{-}Tien Lin},
booktitle = {Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual},
year = {2020}
}
@misc{Pytorch-PCGrad,
author = {Luís Vilaça},
title = {lvilaca16/Pytorch-PCGrad},
url = {https://github.com/lvilaca16/Pytorch-PCGrad},
year = {2026}
}