调用自定义类

概述

在静态图模式下，通过使用jit_class修饰自定义类，用户可以创建、调用该自定义类的实例，并且可以获取其属性和方法。

jit_class应用于静态图模式，扩充完善静态图编译语法的支持范围。在动态图模式即PyNative模式下，jit_class的使用不影响PyNative模式的执行逻辑。

本文档主要介绍jit_class装饰器的使用方法和使用须知，以便您可以更有效地使用jit_class装饰器功能。

jit_class修饰自定义类

使用@jit_class修饰自定义类后，支持创建、调用该自定义类的实例，获取其属性和方法。

[1]:

import numpy as np
import mindspore.nn as nn
import mindspore as ms

@ms.jit_class
class InnerNet:
    value = ms.Tensor(np.array([1, 2, 3]))

class Net(nn.Cell):
    def construct(self):
        return InnerNet().value

ms.set_context(mode=ms.GRAPH_MODE)
net = Net()
out = net()
print(out)

[1 2 3]

jit_class支持自定义类嵌套使用、自定义类与nn.Cell嵌套使用的场景。需要注意的是，类继承时，如果父类使用了jit_class，子类也会具有jit_class的能力。

[2]:

import numpy as np
import mindspore.nn as nn
import mindspore as ms

@ms.jit_class
class Inner:
    def __init__(self):
        self.value = ms.Tensor(np.array([1, 2, 3]))

@ms.jit_class
class InnerNet:
    def __init__(self):
        self.inner = Inner()

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.inner_net = InnerNet()

    def construct(self):
        out = self.inner_net.inner.value
        return out

ms.set_context(mode=ms.GRAPH_MODE)
net = Net()
out = net()
print(out)

[1 2 3]

jit_class仅支持修饰自定义类，不支持nn.Cell和非class类型。执行下面用例，将会出现报错。

import mindspore.nn as nn
import mindspore as ms

@ms.jit_class
class Net(nn.Cell):
    def construct(self, x):
        return x

ms.set_context(mode=ms.GRAPH_MODE)
x = ms.Tensor(1)
net = Net()
net(x)

报错信息如下：

TypeError: Decorator jit_class is used for user-defined classes and cannot be used for nn.Cell: Net<>.

import mindspore as ms

@ms.jit_class
def func(x, y):
    return x + y

func(1, 2)

报错信息如下：

TypeError: Decorator jit_class can only be used for class type, but got <function func at 0x7fee33c005f0>.

获取自定义类的属性和方法

支持通过类名或实例调用类的属性和方法。

[3]:

import mindspore.nn as nn
import mindspore as ms

@ms.jit_class
class InnerNet:
    def __init__(self, val):
        self.number = val

    def act(self, x, y):
        return self.number * (x + y)

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.inner_net = InnerNet(2)

    def construct(self, x, y):
        return self.inner_net.number + self.inner_net.act(x, y)

ms.set_context(mode=ms.GRAPH_MODE)
x = ms.Tensor(2, dtype=ms.int32)
y = ms.Tensor(3, dtype=ms.int32)
net = Net()
out = net(x, y)
print(out)

创建自定义类的实例

静态图模式下，创建自定义类的实例时，参数要求为常量。

[4]:

import numpy as np
import mindspore.nn as nn
import mindspore as ms

@ms.jit_class
class InnerNet:
    def __init__(self, val):
        self.number = val + 3

class Net(nn.Cell):
    def construct(self):
        net = InnerNet(2)
        return net.number

ms.set_context(mode=ms.GRAPH_MODE)
net = Net()
out = net()
print(out)

对于其它场景，创建自定义类的实例时，没有参数必须是常量的限制。例如下面的用例：

[5]:

import numpy as np
import mindspore.nn as nn
import mindspore as ms

@ms.jit_class
class InnerNet:
    def __init__(self, val):
        self.number = val + 3

class Net(nn.Cell):
    def __init__(self, val):
        super(Net, self).__init__()
        self.inner = InnerNet(val)

    def construct(self):
        return self.inner.number

ms.set_context(mode=ms.GRAPH_MODE)
x = ms.Tensor(2, dtype=ms.int32)
net = Net(x)
out = net()
print(out)

调用自定义类的实例

调用自定义类的实例时，将会调用该类的__call__函数方法。

[6]:

import numpy as np
import mindspore.nn as nn
import mindspore as ms

@ms.jit_class
class InnerNet:
    def __init__(self, number):
        self.number = number

    def __call__(self, x, y):
        return self.number * (x + y)

class Net(nn.Cell):
    def construct(self, x, y):
        net = InnerNet(2)
        out = net(x, y)
        return out

ms.set_context(mode=ms.GRAPH_MODE)
x = ms.Tensor(2, dtype=ms.int32)
y = ms.Tensor(3, dtype=ms.int32)
net = Net()
out = net(x, y)
print(out)

如果该类没有定义__call__函数，将会报错提示。执行下面用例，将会出现报错。

import numpy as np
import mindspore.nn as nn
import mindspore as ms

@ms.jit_class
class InnerNet:
    def __init__(self, number):
        self.number = number

class Net(nn.Cell):
    def construct(self, x, y):
        net = InnerNet(2)
        out = net(x, y)
        return out

ms.set_context(mode=ms.GRAPH_MODE)
x = ms.Tensor(2, dtype=ms.int32)
y = ms.Tensor(3, dtype=ms.int32)
net = Net()
out = net(x, y)
print(out)

报错信息如下：

RumtimeError: MsClassObject: 'InnerNet' has no `__call__` function, please check the code.