Example code for fashion-MNIST by tensorflow and pytorch with docker to test server's enviornment.

Dockerfile

FROM docker.io/nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04
RUN apt update -y && apt install -y python3-pip
RUN pip install tensorflow torch torchvision tqdm pillow

fashion_mnist_tensorflow.py

# Modified from https://github.com/tensorflow/docs/blob/master/site/en/tutorials/keras/classification.ipynb
import tensorflow as tf
import numpy as np

# gpu setup
gpus = tf.config.experimental.list_physical_devices("GPU")
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# dataset
fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
train_images = train_images / 255.0
test_images = test_images / 255.0

# model
model = tf.keras.Sequential(
    [
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(16384, activation="relu"),
        tf.keras.layers.Dense(16384, activation="relu"),
        tf.keras.layers.Dense(16384, activation="relu"),
        tf.keras.layers.Dense(16384, activation="relu"),
        tf.keras.layers.Dense(16384, activation="relu"),
        tf.keras.layers.Dense(128, activation="relu"),
        tf.keras.layers.Dense(10),
    ]
)

# model parameters
model.compile(
    optimizer="adam",
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
)
# train
model.fit(train_images, train_labels, batch_size=64, shuffle=True, epochs=30)
# evalaute
test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)
print("\nTest accuracy:", test_acc)

fashion_mnist_torch.py

# Modified from https://github.com/pranay414/Fashion-MNIST-Pytorch/blob/master/fashion_mnist.ipynb
import torch
from torch import nn
from torchvision import datasets, transforms
from tqdm import tqdm

# dataset
transform = transforms.Compose(
    [transforms.ToTensor()]
)
trainset = datasets.FashionMNIST(
    ".keras/datasets/fashion-mnist/", download=True, train=True, transform=transform
)
testset = datasets.FashionMNIST(
    ".keras/datasets/fashion-mnist/", download=True, train=False, transform=transform
)
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True
)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True
)

# model
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 16384),
    nn.ReLU(),
    nn.Linear(16384, 16384),
    nn.ReLU(),
    nn.Linear(16384, 16384),
    nn.ReLU(),
    nn.Linear(16384, 16384),
    nn.ReLU(),
    nn.Linear(16384, 16384),
    nn.ReLU(),
    nn.Linear(16384, 16384),
    nn.ReLU(),
    nn.Linear(16384, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
    nn.LogSoftmax(dim=1),
).cuda()

# model parameters
criterion = nn.NLLLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.002)
epochs = 30

# train
train_losses, test_losses = [], []
for e in range(epochs):
    running_loss = 0
    model.train()
    for images, labels in tqdm(trainloader):
        images = images.cuda()
        labels = labels.cuda()

        # training step
        optimizer.zero_grad()
        output = model.forward(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    train_losses.append(running_loss / len(trainloader))

    # evaluate
    test_loss = 0
    accuracy = 0
    with torch.no_grad():
        model.eval()
        for images, labels in testloader:
            images = images.cuda()
            labels = labels.cuda()

            # evaluating step
            log_ps = model(images)
            test_loss += criterion(log_ps, labels)

            # accuracy step
            ps = torch.exp(log_ps)
            top_p, top_class = ps.topk(1, dim=1)
            equals = top_class == labels.view(*top_class.shape)
            accuracy += torch.mean(equals.type(torch.FloatTensor))
    test_losses.append(test_loss / len(testloader))
    print(
        "Epoch: {:2d}/{:2d}  ".format(e + 1, epochs),
        "Training loss: {:.3f}  ".format(running_loss / len(trainloader)),
        "Test loss: {:.3f}  ".format(test_loss / len(testloader)),
        "Test Accuracy: {:.3f}".format(accuracy / len(testloader)),
    )

test_cnn.py

# CNN model can take more power consumption
# change the model here
# pytorch
model = nn.Sequential(
    nn.Conv2d(1, 2048, 3, padding=1),
    nn.ReLU(),
    nn.Conv2d(2048, 2048, 3, padding=1),
    nn.ReLU(),
    nn.Conv2d(2048, 2048, 3, padding=1),
    nn.ReLU(),
    nn.Conv2d(2048, 2048, 3, padding=1),
    nn.ReLU(),
    nn.Conv2d(2048, 2048, 3, padding=1),
    nn.ReLU(),
    nn.Conv2d(2048, 1, 3, padding=1),
    nn.Flatten(),
    nn.Linear(28 * 28, 10),
    nn.LogSoftmax(dim=1),
).cuda()

# tensorflow
model = tf.keras.Sequential(
    [
        tf.keras.layers.Conv2D(2048, 3, padding="same", activation="relu", input_shape=(1, 28, 28)),
        tf.keras.layers.Conv2D(2048, 3, padding="same", activation="relu"),
        tf.keras.layers.Conv2D(2048, 3, padding="same", activation="relu"),
        tf.keras.layers.Conv2D(2048, 3, padding="same", activation="relu"),
        tf.keras.layers.Conv2D(1,    3, padding="same", activation="relu"),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(10),
    ]
)