75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
import tensorlayer as T
|
|
from dragon.vm.tensorlayer.layers import Dense
|
|
from dragon.vm.tensorlayer.models import Model
|
|
import dragon.vm.tensorlayer as tl
|
|
|
|
import dragon as dg
|
|
import argparse
|
|
import numpy as np
|
|
|
|
X_train, y_train, X_val, y_val, X_test, y_test = T.files.load_mnist_dataset(shape=(-1, 784))
|
|
|
|
|
|
class MLP(Model):
|
|
|
|
def __init__(self):
|
|
super(MLP, self).__init__()
|
|
self.dense1 = Dense(n_units=800, act=tl.act.relu, in_channels=784)
|
|
self.dense2 = Dense(n_units=800, act=tl.act.relu, in_channels=800)
|
|
self.dense3 = Dense(n_units=10, act=tl.act.relu, in_channels=800)
|
|
|
|
def forward(self, x):
|
|
z = self.dense1(x)
|
|
z = self.dense2(z)
|
|
out = self.dense3(z)
|
|
return out
|
|
|
|
|
|
class Classifier(object):
|
|
"""The base classifier class."""
|
|
|
|
# TensorSpec for graph execution
|
|
image_spec = dg.Tensor([None, 3, 32, 32], 'float32')
|
|
label_spec = dg.Tensor([None], 'int64')
|
|
|
|
def __init__(self, optimizer):
|
|
super(Classifier, self).__init__()
|
|
self.net = MLP()
|
|
self.optimizer = optimizer
|
|
self.params = self.net.trainable_weights
|
|
|
|
def step(self, image, label):
|
|
with dg.GradientTape() as tape:
|
|
logit = self.net(image)
|
|
# logit = dg.cast(logit, 'float64')
|
|
logit = dg.cast(dg.math.argmax(logit, -1), 'int32')
|
|
# label = dg.cast(label, 'float32')
|
|
# print("logit :\n", logit, label)
|
|
# loss = dg.losses.smooth_l1_loss([logit, label])
|
|
# loss = tl.losses.sparse_softmax_crossentropy(logit, label)
|
|
loss = dg.math.sum(
|
|
(logit - label) * (logit - label)
|
|
) # dg.losses.sparse_softmax_cross_entropy([logit, label])
|
|
accuracy = dg.math.mean(dg.math.equal([logit, label]).astype('float32'))
|
|
grads = tape.gradient(loss, self.params)
|
|
self.optimizer.apply_gradients(zip(self.params, grads))
|
|
return loss, accuracy, self.optimizer
|
|
|
|
|
|
if __name__ == '__main__':
|
|
dg.autograph.set_execution('EAGER_MODE')
|
|
# Define the model
|
|
model = Classifier(dg.optimizers.SGD(base_lr=0.001, momentum=0.9, weight_decay=1e-4))
|
|
# Main loop
|
|
batch_size = 200
|
|
for i in range(50):
|
|
for X_batch, y_batch in T.iterate.minibatches(X_train, y_train, batch_size, shuffle=True):
|
|
image = dg.EagerTensor(X_batch, copy=False)
|
|
label = dg.EagerTensor(y_batch, copy=False, dtype='float32')
|
|
loss, accuracy, _ = model.step(image, label)
|
|
if i % 20 == 0:
|
|
dg.logging.info(
|
|
'Iteration %d, lr = %s, loss = %.5f, accuracy = %.3f' %
|
|
(i, str(model.optimizer.base_lr), loss, accuracy)
|
|
)
|