main.py

import tensorflow as tf  # TF 2.0
from sn import SpectralNormalization

from tensorflow.keras import datasets, layers, models


batch_size = 200
buffer_size = 2000
num_epochs = 200
method = 'func'

(train_images, train_labels), (_, _) = datasets.mnist.load_data()

train_images = train_images.reshape((60000, 28, 28, 1)).astype('float32')

train_images = train_images / 255.0

train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels)).shuffle(buffer_size).batch(batch_size)

if method == 'func':
    inputs = layers.Input(shape=(28,28,1))
    x = SpectralNormalization(layers.Conv2D(32, (3, 3), activation='relu'))(inputs)
    x = layers.MaxPooling2D((2, 2))(x)
    x = SpectralNormalization(layers.Conv2D(64, (3, 3), activation='relu'))(x)
    x = layers.MaxPooling2D((2, 2))(x)
    x = layers.Flatten()(x)
    x = layers.Dense(64, activation='relu')(x)
    output = layers.Dense(10, activation='softmax')(x)
    model = models.Model(inputs=inputs, outputs=output)
else:
    model = models.Sequential()
    model.add(SpectralNormalization(layers.Conv2D(32, (3, 3), activation='relu')))
    model.add(layers.MaxPooling2D((2, 2)))
    model.add(SpectralNormalization(layers.Conv2D(64, (3, 3), activation='relu')))
    model.add(layers.MaxPooling2D((2, 2)))
    model.add(layers.Flatten())
    model.add(layers.Dense(64, activation='relu'))
    model.add(layers.Dense(10, activation='softmax'))


def loss(model, x, y):
    y_ = model(x)

    return loss_object(y_true=y, y_pred=y_)


def grad(model, inputs, targets):
    with tf.GradientTape() as tape:
      loss_value = loss(model, inputs, targets)
    return loss_value, tape.gradient(loss_value, model.trainable_variables)


loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)

train_loss_results = []
train_accuracy_results = []

for epoch in range(num_epochs):
    epoch_loss_avg = tf.keras.metrics.Mean()
    epoch_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()

    for x, y in train_dataset:
        loss_value, grads = grad(model, x, y)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        epoch_loss_avg(loss_value)
        epoch_accuracy(y, model(x))

    train_loss_results.append(epoch_loss_avg.result())
    train_accuracy_results.append(epoch_accuracy.result())

    print("Epoch {:03d}: Loss: {:.3f}, Acc: {:.3%}".format(epoch,
                                                           epoch_loss_avg.result(),
                                                           epoch_accuracy.result()))