|
- #! /usr/bin/python
- # -*- coding: utf-8 -*-
-
- import os
- # os.environ['TL_BACKEND'] = 'tensorflow'
- os.environ['TL_BACKEND'] = 'mindspore'
-
-
- import tensorlayerx as tlx
- from tensorlayerx.nn import Module
- from tensorlayerx.nn import Linear, Dropout, Conv2d, MaxPool2d, Flatten
- from tensorlayerx.dataflow import Dataset
- from tensorlayerx.backend import BACKEND
- class CNN(Module):
-
- def __init__(self):
- super(CNN, self).__init__()
- W_init = tlx.nn.initializers.truncated_normal(stddev=5e-2)
- W_init2 = tlx.nn.initializers.truncated_normal(stddev=0.04)
- b_init2 = tlx.nn.initializers.constant(value=0.1)
-
- self.conv1 = Conv2d(64, (5, 5), (2, 2), padding='SAME', W_init=W_init, name='conv1', in_channels=3, data_format = 'channels_first')
- self.maxpool1 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool1', data_format = 'channels_first')
-
- self.conv2 = Conv2d(
- 64, (5, 5), (2, 2), padding='SAME', act=tlx.nn.ReLU, W_init=W_init, b_init=None, name='conv2', in_channels=64, data_format = 'channels_first'
- )
- self.maxpool2 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool2', data_format = 'channels_first')
-
- self.flatten = Flatten(name='flatten')
- self.linear1 = Linear(384, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear1', in_features=256)
- self.linear2 = Linear(192, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear2', in_features=384)
- self.linear3 = Linear(10, act=None, W_init=W_init2, name='linear3', in_features=192)
-
- def forward(self, x):
- z = self.conv1(x)
- z = self.maxpool1(z)
- z = self.conv2(z)
- z = self.maxpool2(z)
- z = self.flatten(z)
- z = self.linear1(z)
- z = self.linear2(z)
- z = self.linear3(z)
- return z
-
- cnn = CNN()
- if BACKEND == 'tensorflow':
- cnn.save_standard_weights('tf_model.npz')
- print("TF SAVE DONE!")
- elif BACKEND == 'mindspore':
- cnn.load_standard_weights("tf_model.npz", weights_from = 'tensorflow', weights_to = 'mindspore')
- print("MS LOAD DONE!")
|