自定义autograd function

最新推荐文章于 2025-04-19 16:33:14 发布

seamanj

最新推荐文章于 2025-04-19 16:33:14 发布

阅读量650

点赞数 3

分类专栏： pytorch

本文链接：https://blog.youkuaiyun.com/seamanj/article/details/110083044

版权

pytorch 专栏收录该内容

29 篇文章

订阅专栏

这篇博客详细介绍了如何在PyTorch中为自定义函数SegmentConsensus实现梯度运算。SegmentConsensus根据输入的共识类型（如平均或身份）进行操作，并在前向传播中计算输出。在反向传播时，它根据共识类型计算输入的梯度。此功能对于视频理解等任务中的时空特征聚合至关重要。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在TSN代码中， segmentconsensus是一个自定义函数，所以要写一下它对应的梯度运算


# tj : https://blog.youkuaiyun.com/tsq292978891/article/details/79364140
class SegmentConsensus(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_tensor, consensus_type, dim):


        ctx.consensus_type = consensus_type
        ctx.dim = dim
        ctx.shape = input_tensor.size()
        if consensus_type == 'avg':
            output = input_tensor.mean(dim=dim, keepdim=True)
        elif consensus_type == 'identity':
            output = input_tensor
        else:
            output = None

        return output



    @staticmethod
    def backward(ctx, grad_output):


        if ctx.consensus_type == 'avg':
            grad_in = grad_output.expand(ctx.shape) / float(ctx.shape[ctx.dim])
        elif ctx.consensus_type == 'identity':
            grad_in = grad_output
        else:
            grad_in = None

        return grad_in, None, None
        # tj : 常数要返回 None


# tj : https://blog.youkuaiyun.com/tsq292978891/article/details/79364140
class ConsensusModule(torch.nn.Module):

    def __init__(self, consensus_type, dim=1):
        super(ConsensusModule, self).__init__()
        self.consensus_type = consensus_type if consensus_type != 'rnn' else 'identity'
        self.dim = dim

    def forward(self, input):
        return SegmentConsensus.apply(input, self.consensus_type, self.dim)