Tensorboard自定义网络

keras.Sequential
keara.layers.Layer
keras.Model

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


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])
        self.bias = self.add_variable('b', [output_dim])
    def call(self, inputs, training = None):
        out = inputs @ self.kernel +self.bias
        return out

class MyModel(keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = MyDense(28*28, 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):
        out = self.fc1(inputs)
        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

def preprocess(x, y):
    x= tf.cast(x, dtype = tf.float32) / 255.
    x = tf.reshape(x, [28*28])
    y = tf.cast(y, dtype = tf.int32)
    y = tf.one_hot(y, depth = 10)
    return x, y
(x, y), (x_test, y_test) = keras.datasets.mnist.load_data()
batchsz =128
db =tf.data.Dataset.from_tensor_slices((x, y))
db = db.map(preprocess).shuffle(10000).batch(batchsz)
db_test = tf.data.Dataset.from_tensor_slices((x_test, y_test))
db_test = db_test.map(preprocess).shuffle(10000).batch(batchsz)
network = MyModel()
network.compile(optimizer = optimizers.Adam(lr = 1e-3),
                loss = tf.losses.CategoricalCrossentropy(from_logits = True),
                metrics = ['accuracy'])
network.fit(db, epochs =5, validation_data =db_test,
            validation_freq =2)
network.evaluate(db_test)