Tensorflow cifar10自定义网络实战

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import datasets, layers, optimizers

def preprocess(x, y):
    #[0-255] > [-1, -1]
    x = 2 * tf.cast(x, dtype = tf.float32) / 255. - 1.
    y = tf.cast(y, dtype = tf.int32)
    return x, y
batchz = 128

#[50k, 32,32, 3], [10k,1]
(x, y), (x_test, y_test) = datasets.cifar10.load_data()
y = tf.squeeze(y)
#[10k]
y_test = tf.squeeze(y_test)
y = tf.one_hot(y, depth = 10)
y_test = tf.one_hot(y_test, depth = 10)

train_db = tf.data.Dataset.from_tensor_slices((x, y))
train_db = train_db.map(preprocess).shuffle(10000).batch(batchz)
test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_db = test_db.map(preprocess).shuffle(10000).batch(batchz)

class MyDense(layers.Layer):
    def __init__(self, input_dim, output_dim):
        super(MyDense, self).__init__()
        self.kernel = self.add_variable('w', [input_dim, output_dim])
    def call(self, inputs, training = None):
        out = inputs @ self.kernel
        return out
class MyModel(keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = MyDense(32*32*3, 256)
        self.fc2 = MyDense(256, 128)
        self.fc3 = MyDense(128, 64)
        self.fc4 = MyDense(64, 32)
        self.fc5 = MyDense(32, 10)

    def call(self, inputs, training=None):
        x = tf.reshape(inputs, [-1, 32*32*3])
        out = self.fc1(x)
        out = tf.nn.relu(out)
        out = self.fc2(out)
        out = tf.nn.relu(out)
        out = self.fc3(out)
        out = tf.nn.relu(out)
        out = self.fc4(out)
        out = tf.nn.relu(out)
        out = self.fc5(out)
        return out
network = MyModel()
network.compile(optimizer = optimizers.Adam(lr = 1e-3),
                loss = tf.losses.CategoricalCrossentropy(from_logits = True),
                metrics = ['accuracy'])
network.fit(test_db, epochs = 1, validation_data = test_db, validation_freq =1)
loss,acc = network.evaluate(test_db)
print('loss:', loss, 'accuracy:' ,acc)
network.save_weights('cifar10')
print('network_saved')
del network
network = MyModel()
network.compile(optimizer = optimizers.Adam(lr = 1e-3),
                loss = tf.losses.CategoricalCrossentropy(from_logits = True),
                metrics = ['accuracy'])
network.load_weights('cifar10')
loss,acc = network.evaluate(test_db)
print('loss:', loss, 'accuracy:' ,acc)