大规模人脸分类—allgather操作(1)

pytorch中 all_gather 操作是不进行梯度回传的。在计算图构建中如果需要经过all_gather操作后，仍需要将梯度回传给各个进程中的allgather前的对应变量，则需要重新继承torch.autograd.Function

https://pytorch.org/docs/stable/autograd.html 中对torch.autograd.Function进行了介绍

https://pytorch.org/docs/stable/notes/extending.html#extending-torch-autograd 中举例介绍如何重新实现其子类

下面代码是为了说明all_gather相关特性及如何实现梯度回传.

\(x,y,z\)都是2x2矩阵，其之间关系为\(y=x+2, z=y*y\)

接下来就需要MPI进行进程间数据传递，将z进行汇总到每个进程即all_gather操作。然后将汇总的矩阵进行相乘，然后求均值。

r对y的导数如下:

\(r=0.25({}_{g_0}y_{11}^2*{}_{g_1}y_{11}^2+{}_{g_0}y_{12}^2*{}_{g_1}y_{12}^2+
{}_{g_0}y_{21}^2*{}_{g_1}y_{21}^2+
{}_{g_0}y_{22}^2*{}_{g_1}y_{22}^2)\)

\(\frac{dr}{d{}_{g_0}y}=
\begin{Bmatrix}
0.5{}_{g_0}y_{11}*{}_{g_1}y_{11}^2 & 0.5{}_{g_0}y_{12}*{}_{g_1}y_{12}^2 \\
0.5{}_{g_0}y_{21}*{}_{g_1}y_{21}^2 & 0.5{}_{g_0}y_{22}*{}_{g_1}y_{22}^2)
\end{Bmatrix}\)

gpu0上x值为\(\begin{Bmatrix} 1 & 1 \\1 & 1 \end{Bmatrix}\)，gpu1上x值为\(\begin{Bmatrix} 0 & 0 \\0 & 0 \end{Bmatrix}\).通过公式可以计算出，r关于gpu0上的y的导数为\(\begin{Bmatrix}6 & 6 \\6 & 6\end{Bmatrix}\),r关于gpu1上的y的导数为\(\begin{Bmatrix}9 & 9 \\9 & 9\end{Bmatrix}\)

import os

import torch

from torch import nn

import sys

sys.path.append('./')

import torch.distributed as dist

from torch.autograd import Variable

from utils import GatherLayer

def test():

    #torch.manual_seed(0)

    torch.backends.cudnn.deterministic=True

    torch.backends.cudnn.benchmark=True

    dist.init_process_group(backend="nccl", init_method="env://")

    rank = dist.get_rank()

    local_rank = int(os.environ.get('LOCAL_RANK', 0))

    world_size = dist.get_world_size()

    torch.cuda.set_device(local_rank)

    print('world_size: {}, rank: {}, local_rank: {}'.format(world_size, rank, local_rank))

    if local_rank == 0:

        x = Variable(torch.ones(2, 2), requires_grad=True).cuda()

    else:

        x = Variable(torch.zeros(2, 2), requires_grad=True).cuda()

    y = x + 2

    y.retain_grad()

    z = y * y

    z_gather = [torch.zeros_like(z) for _ in range(world_size)]

    dist.all_gather(z_gather, z)

    #z_gather = GatherLayer.apply(z)

    r = z_gather[0] * z_gather[1]

    out = r.mean()

    out.backward()

    if local_rank == 0:

        print('rank:0', y.grad)

    else:

        print('rank:1', y.grad)

（1）上述述代码中，先使用pytorch中提供的all_gather操作，运行代码会提示错误。错误信息如下：

Traceback (most recent call last):

  File "test/test_all_gather.py", line 46, in <module>

Traceback (most recent call last):

  File "test/test_all_gather.py", line 46, in <module>

    test()

  File "test/test_all_gather.py", line 36, in test

    out.backward()

  File "/usr/local/lib/python3.6/dist-packages/torch/tensor.py", line 185, in backward

    torch.autograd.backward(self, gradient, retain_graph, create_graph)

  File "/usr/local/lib/python3.6/dist-packages/torch/autograd/__init__.py", line 127, in backward

    allow_unreachable=True)  # allow_unreachable flag

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

    test()

（2）参考https://github.com/Spijkervet/SimCLR/blob/master/simclr/modules/gather.py, 该函数就是继承torch.autograd.Function，实现了all_gather后，梯度也能回传。

上述代码，启用z_gather = GatherLayer.apply(z),就实现了梯度回传功能，打印对变量y的梯度

world_size: 2, rank: 0, local_rank: 0

world_size: 2, rank: 1, local_rank: 1

rank:0 tensor([[6., 6.],

        [6., 6.]], device='cuda:0')

rank:1 tensor([[9., 9.],

        [9., 9.]], device='cuda:1')

GatherLayer类实现如下：

class GatherLayer(torch.autograd.Function):

    """Gather tensors from all process, supporting backward propagation."""

    @staticmethod

    def forward(ctx, input):

        ctx.save_for_backward(input)

        output = [torch.zeros_like(input) for _ in range(dist.get_world_size())]

        dist.all_gather(output, input)

        return tuple(output)

    @staticmethod

    def backward(ctx, *grads):

        (input,) = ctx.saved_tensors

        grad_out = torch.zeros_like(input)

        grad_out[:] = grads[dist.get_rank()]

        return grad_out

下面网址有关all gather梯度传播的讨论

https://discuss.pytorch.org/t/will-dist-all-gather-break-the-auto-gradient-graph/47350

巴特西

大规模人脸分类—allgather操作(1)

最新文章

热门文章