Python源码构造网络约束

概述

MindSpore完成从用户源码到计算图的编译，用户源码基于Python语法编写，当前MindSpore支持将普通函数或者继承自nn.Cell的实例转换生成计算图，暂不支持将任意Python源码转换成计算图，所以对于用户源码支持的写法有所限制，主要包括语法约束和网络定义约束两方面。随着MindSpore的演进，这些约束可能会发生变化。

语法约束

支持的Python数据类型

• Number：包括int、float、bool，不支持复数类型。
• String
• List：当前只支持append方法；List的更新会拷贝生成新的List。
• Tuple
• Dictionary：当前key只支持String类型

MindSpore扩展数据类型

• Tensor：Tensor变量必须是已定义实例。

表达式类型

语句类型

系统函数

• len
• partial
• map
• zip
• range

函数参数

• 参数默认值：目前不支持默认值设为Tensor类型数据，支持int、float、bool、None、str、tuple、list、dict类型数据。
• 可变参数：支持带可变参数网络的推理和训练。
• 键值对参数：目前不支持带键值对参数的函数求反向。
• 可变键值对参数：目前不支持带可变键值对的函数求反向。

操作符

索引操作

索引操作包含tuple和Tensor的索引操作。下面重点介绍一下Tensor的索引取值和赋值操作，取值以tensor_x[index]为例，赋值以tensor_x[index] = u为例进行详细说明。其中tensor_x是一个Tensor，对其进行切片操作；index表示索引，u表示赋予的值，可以是scalar或者Tensor(size=1)。索引类型如下：

• 切片索引：index为slice

  • 取值：tensor_x[start:stop:step]，其中Slice(start:stop:step)与Python的语法相同，这里不再赘述。
  • 赋值：tensor_x[start:stop:step]=u。

• Ellipsis索引：index为ellipsis

  • 取值：tensor_x[...]。
  • 赋值：tensor_x[...]=u。

• 布尔常量索引：index为True，index为False暂不支持。

  • 取值：tensor_x[True]。
  • 赋值：暂不支持。

• Tensor索引：index为Tensor

  • 取值：tensor_x[index]，index必须是int32、int64类型的Tensor，元素取值范围在[0, tensor_x.shape[0])。
  • 赋值：tensor_x[index]=U。
    • tensor_x的数据类型必须是下面一种： float16，float32，int8，uint8。
    • index必须是int32类型的Tensor，元素取值范围在[0, tensor_x.shape[0])。
    • U可以是Number，Tensor，只包含Number的Tuple，只包含Tensor的Tuple。
      • 单个Number和Tuple里的每个Number必须与tensor_x的数据类型属于同一类，即 当tensor_x的数据类型是uint8或者int8时，Number类型应该是int； 当tensor_x的数据类型是float16或者float32时，Number类型应该是float。
      • 单个Tensor和Tuple里的每个Tensor必须与tensor_x的数据类型一致， 单个Tensor时，其shape需等于或者可广播为index.shape + tensor_x.shape[1:]。
      • 包含Number的Tuple需满足下面条件： len(Tuple) = (index.shape + tensor_x.shape[1:])[-1]。
      • 包含Tensor的Tuple需满足下面条件： 每个Tensor的shape一样； (len(Tuple),) + Tensor.shape等于或者可广播为index.shape + tensor_x.shape[1:]。

• None常量索引：index为None

  • 取值：tensor_x[None]，结果与numpy保持一致。
  • 赋值：暂不支持。

• tuple索引：index为tuple

  • tuple元素为slice:
    • 取值：例如tensor_x[::, :4, 3:0:-1]。
    • 赋值：例如tensor_x[::, :4, 3:0:-1]=u。
  • tuple元素为Number:
    • 取值：例如tensor_x[2,1]。
    • 赋值：例如tensor_x[1,4]=u。
  • tuple元素为slice和ellipsis混合情况:
    • 取值：例如tensor_x[..., ::, 1:]
    • 赋值：例如tensor_x[..., ::, 1:]=u
  • 其他情况暂不支持

另外tuple也支持切片取值操作，tuple_x[start:stop:step]，与Python的效果相同，这里不再赘述。

不支持的语法

目前在网络构造函数里面暂不支持以下语法：
raise、 yield、 async for、 with、 async with、 assert、 import、 await。

网络定义约束

整网实例类型

• 带@ms_function装饰器的普通Python函数。
• 继承自nn.Cell的Cell子类。

网络输入类型

• 整网的训练数据输入参数只能是Tensor类型。
• 生成的ANF图里面不能包含这几种常量节点：字符串类型常量、带有Tuple嵌套的常量、带有List嵌套的常量。

网络图优化

在ME前端图优化过程中，会将DataClass类型、Dictionary、List、键值对操作转换为Tuple相关操作。

网络构造组件

其他约束

整网construct函数输入的参数以及使用ms_function装饰器修饰的函数的参数在图编译过程中会进行泛化，不能作为常量输入传给算子使用，如下例所示：

• 错误的写法如下：

  class ExpandDimsTest(Cell):
      def __init__(self):
          super(ExpandDimsTest, self).__init__()
          self.expandDims = P.ExpandDims()
  
      def construct(self, input_x, input_axis):
          return self.expandDims(input_x, input_axis)
  expand_dim = ExpandDimsTest()
  input_x = Tensor(np.random.randn(2,2,2,2).astype(np.float32))
  expand_dim(input_x, 0)

  在示例中，ExpandDimsTest是一个只有单算子的网络，网络的输入有input_x和input_axis两个。因为ExpandDims算子的第二个输入需要是常量，这是因为在图编译过程中推导ExpandDims算子输出维度的时候需要用到，而input_axis作为网络参数输入会泛化成变量，无法确定其值，从而无法推导算子的输出维度导致图编译失败。所以在图编译阶段需要值推导的输入都应该是常量输入。在API中，这类算子需要常量输入的参数会进行说明，标注"constant input is needed"。

• 正确的写法是在construct函数里面对算子的常量输入直接填入需要的值或者是一个类的成员变量，如下：

  class ExpandDimsTest(Cell):
      def __init__(self, axis):
          super(ExpandDimsTest, self).__init__()
          self.expandDims = P.ExpandDims()
          self.axis = axis
  
      def construct(self, input_x):
          return self.expandDims(input_x, self.axis)
  axis = 0
  expand_dim = ExpandDimsTest(axis)
  input_x = Tensor(np.random.randn(2,2,2,2).astype(np.float32))
  expand_dim(input_x)