|
- import os
- os.environ['TL_BACKEND'] = 'mindspore'
-
- import mindspore.ops.operations as P
- from mindspore.ops import composite as C
- from mindspore import ParameterTuple
- from mindspore.nn import Momentum, WithLossCell
-
- import numpy as np
- import tensorlayerx as tlx
- import mindspore as ms
- import time
- from tensorlayerx.nn import Module
- from tensorlayerx.nn import Linear
- import mindspore.nn as nn
- from tensorlayerx.dataflow import Dataset, DataLoader
-
- X_train, y_train, X_val, y_val, X_test, y_test = tlx.files.load_mnist_dataset(shape=(-1, 784))
-
-
- class mnistdataset(Dataset):
-
- def __init__(self, data=X_train, label=y_train):
- self.data = data
- self.label = label
-
- def __getitem__(self, index):
- data = self.data[index].astype('float32')
- label = self.label[index].astype('int64')
- return data, label
-
- def __len__(self):
- return len(self.data)
-
- class MLP(Module):
-
- def __init__(self):
- super(MLP, self).__init__()
- self.linear1 = Linear(out_features=800, act=tlx.nn.ReLU, in_features=784)
- self.linear2 = Linear(out_features=800, act=tlx.nn.ReLU, in_features=800)
- self.linear3 = Linear(out_features=10, act=tlx.nn.ReLU, in_features=800)
-
- def forward(self, x):
- z = self.linear1(x)
- z = self.linear2(z)
- out = self.linear3(z)
- return out
-
-
- class GradWrap(Module):
- """ GradWrap definition """
-
- def __init__(self, network):
- super(GradWrap, self).__init__(auto_prefix=False)
- self.network = network
- self.weights = ParameterTuple(filter(lambda x: x.requires_grad, network.get_parameters()))
-
- def forward(self, x, label):
- return C.GradOperation(get_by_list=True)(self.network, self.weights)(x, label)
-
-
-
- net = MLP()
- train_weights = list(filter(lambda x: x.requires_grad, net.get_parameters()))
- optimizer = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.15, 0.8)
-
- criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
- net_with_criterion = WithLossCell(net, criterion)
- train_network = GradWrap(net_with_criterion)
- train_network.set_train()
-
- train_dataset = mnistdataset(data=X_train, label=y_train)
- train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
-
- n_epoch = 50
-
- for epoch in range(n_epoch):
- start_time = time.time()
- train_network.set_train()
- train_loss, train_acc, n_iter = 0, 0, 0
- for X_batch, y_batch in train_loader:
- output = net(X_batch)
- loss_output = criterion(output, y_batch)
- grads = train_network(X_batch, y_batch)
- success = optimizer(grads)
- loss = loss_output.asnumpy()
- train_loss += loss
- n_iter += 1
- train_acc += np.mean((P.Equal()(P.Argmax(axis=1)(output), y_batch).asnumpy()))
- print("Epoch {} of {} step {} took {}".format(epoch + 1, n_epoch, n_iter, time.time() - start_time))
- print(" train loss: {}".format(train_loss / n_iter))
- print(" train acc: {}".format(train_acc / n_iter))
|