!67 add cpu st lenet

Merge pull request !67 from kisnwang/add-cpu-st

tests/st/networks/test_cpu_lenet.py

@@ -12,25 +12,44 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""
-Function: 
-    test network
-Usage: 
-    python test_network_main.py --net lenet --target Davinci
-"""
-import os
-import time
 import pytest
-import numpy as np
-import argparse
-import mindspore.nn as nn
-from mindspore.common.tensor import Tensor
 from mindspore.nn import TrainOneStepCell, WithLossCell
 import mindspore.context as context
 from mindspore.nn.optim import Momentum
-from models.lenet import LeNet
+import numpy as np
-from models.resnetv1_5 import resnet50
+import mindspore.nn as nn
-from models.alexnet import AlexNet
+from mindspore.ops import operations as P
+from mindspore import Tensor
+
+class LeNet(nn.Cell):
+    def __init__(self):
+        super(LeNet, self).__init__()
+        self.relu = P.ReLU()
+        self.batch_size = 32
+
+        self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=0, has_bias=False, pad_mode='valid')
+        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0, has_bias=False, pad_mode='valid')
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.reshape = P.Reshape()
+        self.fc1 = nn.Dense(400, 120)
+        self.fc2 = nn.Dense(120, 84)
+        self.fc3 = nn.Dense(84, 10)
+
+    def construct(self, input_x):
+        output = self.conv1(input_x)
+        output = self.relu(output)
+        output = self.pool(output)
+        output = self.conv2(output)
+        output = self.relu(output)
+        output = self.pool(output)
+        output = self.reshape(output, (self.batch_size, -1))
+        output = self.fc1(output)
+        output = self.relu(output)
+        output = self.fc2(output)
+        output = self.relu(output)
+        output = self.fc3(output)
+        return output
+
 context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
 
 def train(net, data, label):
@@ -48,15 +67,6 @@ def train(net, data, label):
     print("+++++++++++++++++++++++++++")
     assert res
 
-@pytest.mark.level0
-@pytest.mark.platform_x86_cpu
-@pytest.mark.env_onecard
-def test_resnet50():
-    data = Tensor(np.ones([32, 3 ,224, 224]).astype(np.float32) * 0.01)
-    label = Tensor(np.ones([32]).astype(np.int32))
-    net = resnet50(32, 10)
-    train(net, data, label)
-
 @pytest.mark.level0
 @pytest.mark.platform_x86_cpu
 @pytest.mark.env_onecard
@@ -65,12 +75,3 @@ def test_lenet():
     label = Tensor(np.ones([32]).astype(np.int32))
     net = LeNet()
     train(net, data, label)
-
-@pytest.mark.level0
-@pytest.mark.platform_x86_cpu
-@pytest.mark.env_onecard
-def test_alexnet():
-    data = Tensor(np.ones([32, 3 ,227, 227]).astype(np.float32) * 0.01)
-    label = Tensor(np.ones([32]).astype(np.int32))
-    net = AlexNet()
-    train(net, data, label)