class DenseUnit(tf.keras.Model):
def __init__(self, filter_out, kernel_size):
super(DenseUnit, self).__init__()
self.bn = tf.keras.layers.BatchNormalization()
self.conv = tf.keras.layers.Conv2D(filter_out, kernel_size, padding='same')
self.concat = tf.keras.layers.Concatenate()
def call(self, x, training=False, mask=None): # x : (Batch, H, W, Ch_in)
h = self.bn(x, training=training)
h = tf.nn.relu(h)
h = self.conv(h) # h : (Batch, H, W, filter_output)
return self.concat([x, h]) # (Batch, H, W, (Ch_in + filter_output))
class DenseLayer(tf.keras.Model):
def __init__(self, num_unit, growth_rate, kernel_size):
super(DenseLayer, self).__init__()
self.sequence = list()
for idx in range(num_unit):
self.sequence.append(DenseUnit(growth_rate, kernel_size))
def call(self, x, training=False, mask=None):
for unit in self.sequence:
x = unit(x, training=training)
return x
class TransitionLayer(tf.keras.Model):
def __init__(self, filters, kernel_size):
super(TransitionLayer, self).__init__()
self.conv = tf.keras.layers.Conv2D(filters, kernel_size, padding='same')
self.pool = tf.keras.layers.MaxPool2D()
def call(self, x, training=False, mask=None):
x = self.conv(x)
return self.pool(x)
class DenseNet(tf.keras.Model):
def __init__(self):
super(DenseNet, self).__init__()
self.conv1 = tf.keras.layers.Conv2D(8, (3, 3), padding='same', activation='relu') # 28 x 28 x 8
self.dl1 = DenseLayer(2, 4, (3, 3)) # 28 x 28 x 16
self.tr1 = TransitionLayer(16, (3, 3)) # 14 x 14 x 16
self.dl2 = DenseLayer(2, 8, (3, 3)) # 14 x 14 x 32
self.tr2 = TransitionLayer(32, (3, 3)) # 7 x 7 x 32
self.dl3 = DenseLayer(2, 16, (3, 3)) # 7 x 7 x 64
self.flatten = tf.keras.layers.Flatten()
self.dense1 = tf.keras.layers.Dense(128, activation='relu')
self.dense2 = tf.keras.layers.Dense(10 , activation='softmax')
def call(self, x, training=False, mask=None):
x = self.conv1(x)
x = self.dl1(x, training=training)
x = self.tr1(x)
x = self.dl2(x, training=training)
x = self.tr2(x)
x = self.dl3(x, training=training)
x = self.flatten(x)
x = self.dense1(x)
return self.dense2(x)
# Implement training loop
@tf.function
def train_step(model, images, labels, loss_object, optimizer, train_loss, train_accuracy):
with tf.GradientTape() as tape:
predictions = model(images, training=True)
loss = loss_object(labels, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss(loss)
train_accuracy(labels, predictions)
# Implement algorithm test
@tf.function
def test_step(model, images, labels, loss_object, test_loss, test_accuracy):
predictions = model(images, training=False)
t_loss = loss_object(labels, predictions)
test_loss(t_loss)
test_accuracy(labels, predictions)
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
x_train = x_train[..., tf.newaxis].astype(np.float32)
x_test = x_test[..., tf.newaxis].astype(np.float32)
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)
for epoch in range(EPOCHS):
for images, labels in train_ds:
train_step(model, images, labels, loss_object, optimizer, train_loss, train_accuracy)
for test_images, test_labels in test_ds:
test_step(model, test_images, test_labels, loss_object, test_loss, test_accuracy)
template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}'
print(template.format(epoch + 1,
train_loss.result(),
train_accuracy.result() * 100,
test_loss.result(),
test_accuracy.result() * 100))
train_loss.reset_states()
train_accuracy.reset_states()
test_loss.reset_states()
test_accuracy.reset_states()
Epoch 1, Loss: 0.11307405680418015, Accuracy: 96.75333404541016, Test Loss: 0.05511000007390976, Test Accuracy: 98.29000091552734
Epoch 2, Loss: 0.057570938020944595, Accuracy: 98.4383316040039, Test Loss: 0.0530952550470829, Test Accuracy: 98.3499984741211
Epoch 3, Loss: 0.043755121529102325, Accuracy: 98.76333618164062, Test Loss: 0.05816405266523361, Test Accuracy: 98.44999694824219
Epoch 4, Loss: 0.04333307594060898, Accuracy: 98.77333068847656, Test Loss: 0.0660046860575676, Test Accuracy: 98.61000061035156
Epoch 5, Loss: 0.030361121520400047, Accuracy: 99.1500015258789, Test Loss: 0.12039968371391296, Test Accuracy: 97.43000030517578
Epoch 6, Loss: 0.03443685546517372, Accuracy: 99.04999542236328, Test Loss: 0.04854564368724823, Test Accuracy: 98.68000030517578
Epoch 7, Loss: 0.03254585340619087, Accuracy: 99.19833374023438, Test Loss: 0.054437797516584396, Test Accuracy: 98.8699951171875
Epoch 8, Loss: 0.02431321144104004, Accuracy: 99.3116683959961, Test Loss: 0.06114092096686363, Test Accuracy: 98.6199951171875
Epoch 9, Loss: 0.027243511751294136, Accuracy: 99.29833221435547, Test Loss: 0.07321697473526001, Test Accuracy: 98.62999725341797
Epoch 10, Loss: 0.023852670565247536, Accuracy: 99.40166473388672, Test Loss: 0.06338458508253098, Test Accuracy: 98.86000061035156