随机数生成函数导致模型速度越来越慢

[huan666]

1 系统环境

硬件环境(Ascend/GPU/CPU): Ascend/GPU/CPU

MindSpore版本: mindspore=2.0.0

执行模式（PyNative/ Graph）:不限

Python版本: Python=3.9.13

操作系统平台: 不限

2 报错信息

2.1 问题描述

随机数生成函数导致模型速度越来越慢

2.2 脚本代码（代码格式，可上传附件）

import os
import time
import random
import numpy as np
import mindspore as ms
from mindspore import nn, Tensor, ops

seed = 1
ms.set_seed(seed)
np.random.seed(seed)  # Numpy module.
random.seed(seed)  # Python random module.	

class net(nn.Cell):
    def __init__(self, in_dims=4, out_dims=1):
        super().__init__()
        self.fc1 = nn.Dense(in_dims, 4, has_bias=False)
        self.relu = nn.ReLU()
        self.fc2 = nn.Dense(4, out_dims, has_bias=False)

    def construct(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        rand_tensor = ms.numpy.rand(x.shape)
        return x


class NetWithLoss(nn.Cell):
    def __init__(self, net: nn.Cell):
        super(NetWithLoss, self).__init__()
        self.net = net

    def construct(self, state):
        out = self.net(state)
        return out.mean()

net = net()
optimizer = nn.Adam(net.get_parameters(), 0.001) # 使用0学习率，便于观察梯度
net_with_loss = NetWithLoss(net)
net_train_one_step = nn.TrainOneStepCell(net_with_loss, optimizer)

data = ms.Tensor(
    np.random.normal(loc=0.0, scale=1.0, size=(1,4)),
    dtype=ms.float32)

for i in range(1, 1000):
    t1 = time.time()
    out = net_train_one_step(data)
    t2 = time.time()
    print('%d-iter, time:'%(i), t2-t1)

3 根因分析

在PyNative模式下，存在算子缓存机制——每次执行一个算子，如果算子attr和input维度相同，则复用之前的缓存；否则创建新的缓存。

当前脚本中，每次执行rand_tensor = ms.numpy.rand(x.shape)，UniformReal算子的随机数种子会发生变化，因此会生成新的缓存

随着调用次数增加，缓存数量增加，缓存查询操作变慢，导致执行速度变慢

def rand(*shape, dtype=mstype.float32):  
    ...  
    seed = get_seed()  
    if seed is not None:  
        seed1, seed2 = _get_seed(seed, "UniformReal")  
        uniformreal = P.UniformReal(seed=seed1, seed2=seed2)  
    else:  
        uniformreal = P.UniformReal()  
    return uniformreal(size).astype(dtype)

4 解决方案

方案1：使用图模式执行脚本

import os  
import time  
import random  
import numpy as np  
import mindspore as ms  
from mindspore import nn, Tensor, ops, context  
  
# 使用图模式执行  
context.set_context(mode=context.GRAPH_MODE)  
  
seed = 1  
ms.set_seed(seed)  
np.random.seed(seed)  # Numpy module.  
random.seed(seed)  # Python random module.	  
  
class net(nn.Cell):  
    def __init__(self, in_dims=4, out_dims=1):  
        super().__init__()  
        self.fc1 = nn.Dense(in_dims, 4, has_bias=False)  
        self.relu = nn.ReLU()  
        self.fc2 = nn.Dense(4, out_dims, has_bias=False)  
  
    def construct(self, x):  
        x = self.fc1(x)  
        x = self.relu(x)  
        x = self.fc2(x)  
        rand_tensor = ms.numpy.rand(x.shape)  
        return x  
  
  
class NetWithLoss(nn.Cell):  
    def __init__(self, net: nn.Cell):  
        super(NetWithLoss, self).__init__()  
        self.net = net  
  
    def construct(self, state):  
        out = self.net(state)  
        return out.mean()  
  
net = net()  
optimizer = nn.Adam(net.get_parameters(), 0.001) # 使用0学习率，便于观察梯度  
net_with_loss = NetWithLoss(net)  
net_train_one_step = nn.TrainOneStepCell(net_with_loss, optimizer)  
  
data = ms.Tensor(  
    np.random.normal(loc=0.0, scale=1.0, size=(1,4)),  
    dtype=ms.float32)  
  
for i in range(1, 1000):  
    t1 = time.time()  
    out = net_train_one_step(data)  
    t2 = time.time()  
    print('%d-iter, time:'%(i), t2-t1)

方案2：不再随机数种子，避免生成缓存

import os  
import time  
import random  
import numpy as np  
import mindspore as ms  
from mindspore import nn, Tensor, ops  
  
seed = 1  
# 不再设置随机数种子，避免每次执行时，UniformReal算子随机数种子发生变化  
# ms.set_seed(seed)  
np.random.seed(seed)  # Numpy module.  
random.seed(seed)  # Python random module.	  
  
class net(nn.Cell):  
    def __init__(self, in_dims=4, out_dims=1):  
        super().__init__()  
        self.fc1 = nn.Dense(in_dims, 4, has_bias=False)  
        self.relu = nn.ReLU()  
        self.fc2 = nn.Dense(4, out_dims, has_bias=False)  
  
    def construct(self, x):  
        x = self.fc1(x)  
        x = self.relu(x)  
        x = self.fc2(x)  
        rand_tensor = ms.numpy.rand(x.shape)  
        return x  
  
  
class NetWithLoss(nn.Cell):  
    def __init__(self, net: nn.Cell):  
        super(NetWithLoss, self).__init__()  
        self.net = net  
  
    def construct(self, state):  
        out = self.net(state)  
        return out.mean()  
  
net = net()  
optimizer = nn.Adam(net.get_parameters(), 0.001) # 使用0学习率，便于观察梯度  
net_with_loss = NetWithLoss(net)  
net_train_one_step = nn.TrainOneStepCell(net_with_loss, optimizer)  
  
data = ms.Tensor(  
    np.random.normal(loc=0.0, scale=1.0, size=(1,4)),  
    dtype=ms.float32)  
  
for i in range(1, 1000):  
    t1 = time.time()  
    out = net_train_one_step(data)  
    t2 = time.time()  
    print('%d-iter, time:'%(i), t2-t1)