【pytorch->mindspore】1.自定义算子迁移

要迁移的项目为图像压缩算法https://github.com/ywz978020607/HESIC
1.自定义算子迁移--LowerBoundFunction类
为了能够准确迁移底层封装的类,需要详细测试原版类以及迁移测试
pytorch中自定义的算子有torch.autograd.Function

import torch
import torch.nn as nn


class LowerBoundFunction(torch.autograd.Function):
    """Autograd function for the `LowerBound` operator.
    """
    @staticmethod
    def forward(ctx, input_, bound):
        ctx.save_for_backward(input_, bound)
        return torch.max(input_, bound)

    @staticmethod
    def backward(ctx, grad_output):
        input_, bound = ctx.saved_tensors
        pass_through_if = (input_ >= bound) | (grad_output < 0)
        print(pass_through_if)
        print(pass_through_if.type(grad_output.dtype) * grad_output)
        return pass_through_if.type(grad_output.dtype) * grad_output, None

if __name__=="__main__":
    a = torch.Tensor([1,2,3])
    b = torch.Tensor([0,1,5])
    a.requires_grad_(True)
    b.requires_grad_(True)
    c = a*b

    m = LowerBoundFunction.apply(a,b)
    m.backward(c)

通过两行print测试后发现,这个类用于阻断梯度,有点类似Relu的感觉
而mindspore的自定义算子在昇腾、GPU、CPU下定义不同且过于复杂,咨询hw工程师后,准备继承nn.Cell并重载bprop函数实现,测试bprop反向梯度传播如下

# https://gitee.com/mindspore/mindspore/blob/master/tests/ut/python/pynative_mode/test_hook.py#
import numpy as np
import pytest

import mindspore.nn as nn
import mindspore.ops.operations as P
from mindspore import context, Tensor, ParameterTuple
from mindspore.common.initializer import TruncatedNormal
from mindspore.nn import WithLossCell, Momentum
from mindspore.ops import composite as C

context.set_context(mode=context.PYNATIVE_MODE, device_target="CPU")

grad_all = C.GradOperation(get_all=True)
bprop_debug = False

class MulAdd(nn.Cell):
    def __init__(self):
        super(MulAdd, self).__init__()

    def construct(self, x, y):
        return 2 * x * x + y * y

    def bprop(self, x, y, out, dout):
        global bprop_debug
        bprop_debug = True
        return dout, 2 * y


def test_custom_bprop():
    mul_add = MulAdd()
    mul_add.bprop_debug = True
    x = Tensor(np.array([1, 2, 3]).astype(np.int32))
    y = Tensor(np.array([2, 3, 4]).astype(np.int32))
    ret = grad_all(mul_add)(x, y)
    print(ret) #(Tensor(shape=[3], dtype=Int32, value= [1, 1, 1]), Tensor(shape=[3], dtype=Int32, value= [4, 6, 8]))
    assert bprop_debug

##############
#ywz
test_custom_bprop()
print(bprop_debug)
上一篇:编程干货|Python最实用的四个处理图形场景


下一篇:亚麻4年老员工跑毒回国,面字节惨遭秒拒!