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)