I have 4 human face emotion classifier models implenet using mobilenet,resnet,shufflenet,squeezee net stored using .pth now I want to deploy this model using docker such that it takes images input and gives prediction output for different model which will be chosen by user (give steps with code to implement this) 

the model is trained like this:

from google.colab import drive
drive.mount('/content/drive')

# Load the dataset from the zip file
import zipfile
with zipfile.ZipFile('/content/drive/MyDrive/archive (1).zip', 'r') as zip_ref:
    zip_ref.extractall('/content/dataset')

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, ToTensor, Normalize, RandomRotation, RandomHorizontalFlip
from torchvision.datasets import ImageFolder

from matplotlib import pyplot as plt
import numpy as np
import os
import random


# Parameters
IMG_HEIGHT = 48
IMG_WIDTH = 48
batch_size = 32
epochs = 50
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Directories
import torchvision
train_data_dir = '/content/dataset/train'
validation_data_dir = '/content/dataset/test'

mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]

# Define the data transforms for train and test sets
data_transforms = {
    'train': torchvision.transforms.Compose([
        torchvision.transforms.Resize(256),  # Resize the image to 256x256 pixels
        torchvision.transforms.RandomCrop(224),  # Crop a random 224x224 patch from the image
        torchvision.transforms.RandomRotation(30),  # Rotate the image randomly by up to 30 degrees
        torchvision.transforms.RandomHorizontalFlip(),  # Flip the image horizontally with a probability of 0.5
        torchvision.transforms.ToTensor(),  # Convert the image to a PyTorch tensor
        torchvision.transforms.Normalize(mean, std)  # Normalize the image using the mean and standard deviation
    ]),
    'test': torchvision.transforms.Compose([
        torchvision.transforms.Resize(256),  # Resize the image to 256x256 pixels
        torchvision.transforms.CenterCrop(224),  # Crop the center 224x224 patch from the image
        torchvision.transforms.ToTensor(),  # Convert the image to a PyTorch tensor
        torchvision.transforms.Normalize(mean, std)  # Normalize the image using the mean and standard deviation
    ])
}

# Datasets
train_dataset = ImageFolder(train_data_dir, data_transforms['train'])
test_dataset = ImageFolder(validation_data_dir, data_transforms['test'])

# DataLoaders
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

class_labels = ['Angry', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad', 'Surprise']

import matplotlib.pyplot as plt
import numpy as np

batch = next(iter(train_loader))
images, labels = batch

images = images.numpy() / 2 + 0.5

fig, axes = plt.subplots(nrows=3, ncols=3, figsize=(8, 8))

for i, ax in enumerate(axes.flat):
    ax.imshow(np.transpose(images[i], (1, 2, 0)))
    ax.set_title(f"Label: {class_labels[labels[i]]}")

plt.show()


#Custom model

import torch.nn.functional as F
# Model architecture
class EmotionModel(nn.Module):
    def __init__(self):
        super(EmotionModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.pool2 = nn.MaxPool2d(2)
        self.drop2 = nn.Dropout(0.1)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3)
        self.pool3 = nn.MaxPool2d(2)
        self.drop3 = nn.Dropout(0.1)
        self.conv4 = nn.Conv2d(128, 256, kernel_size=3)
        self.pool4 = nn.MaxPool2d(2)
        self.drop4 = nn.Dropout(0.1)
        self.fc1 = nn.Linear(4096, 512)
        self.drop5 = nn.Dropout(0.2)
        self.fc2 = nn.Linear(512, 7)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.drop2(self.pool2(F.relu(self.conv2(x))))
        x = self.drop3(self.pool3(F.relu(self.conv3(x))))
        x = self.drop4(self.pool4(F.relu(self.conv4(x))))
        # print(x.size())  # Add this line to print the size of the tensor
        x = x.view(-1, 4096)
        x = F.relu(self.fc1(x))
        x = self.drop5(x)
        x = self.fc2(x)
        return x



model = EmotionModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())

# Train the model
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(train_loader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}")


# Save the model
torch.save(model.state_dict(), 'emotion_detection_model_50epochs.pth')

# Test the model
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy: {100 * correct / total:.2f}%")


# Plot confusion matrix
from sklearn.metrics import confusion_matrix
import seaborn as sns

y_true = []
y_pred = []
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        y_true.extend(labels.cpu().numpy())
        y_pred.extend(predicted.cpu().numpy())

cm = confusion_matrix(y_true, y_pred)
class_labels = ['Angry', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad', 'Surprise']
sns.heatmap(cm, annot=True, xticklabels=class_labels, yticklabels=class_labels)
plt.show()

# Display a few test images with original and predicted labels
n_display = 6
fig, axes = plt.subplots(1, n_display, figsize=(15, 3))
for i in range(n_display):
    index = random.randint(0, len(test_dataset))
    image, label = test_dataset[index]
    image = image.to(device)
    output = model(image.unsqueeze(0))
    _, prediction = torch.max(output.data, 1)
    orig_label = class_labels[label]
    pred_label = class_labels[prediction.item()]
    
    img = image.cpu().numpy().transpose((1, 2, 0))
    img = img * 0.5 + 0.5  # Un-normalize
    img = np.clip(img, 0, 1)
    axes[i].imshow(img, cmap="gray")
    axes[i].set_title(f"Original: {orig_label}\nPredicted: {pred_label}")
    axes[i].axis("off")
plt.show()

#ResNet

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.models import resnet18

# Define the pre-trained ResNet18 model
model = resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 7)

# Move the model to the GPU if available
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())


# Train the model
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(train_loader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}")

# Save the model
torch.save(model.state_dict(), 'emotion_detection_model_resnet.pth')

model = resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 7)


# Load the pre-trained weights from a file
weights_path = "emotion_detection_model_resnet.pth"
model.load_state_dict(torch.load(weights_path))
model=model.to(device)


# Test the model
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy: {100 * correct / total:.2f}%")


# Plot confusion matrix
from sklearn.metrics import confusion_matrix
import seaborn as sns

y_true = []
y_pred = []
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        y_true.extend(labels.cpu().numpy())
        y_pred.extend(predicted.cpu().numpy())

cm = confusion_matrix(y_true, y_pred)
class_labels = ['Angry', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad', 'Surprise']
sns.heatmap(cm, annot=True, xticklabels=class_labels, yticklabels=class_labels)
plt.show()

# Display a few test images with original and predicted labels
n_display = 6
fig, axes = plt.subplots(1, n_display, figsize=(15, 3))
for i in range(n_display):
    index = random.randint(0, len(test_dataset))
    image, label = test_dataset[index]
    image = image.to(device)
    output = model(image.unsqueeze(0))
    _, prediction = torch.max(output.data, 1)
    orig_label = class_labels[label]
    pred_label = class_labels[prediction.item()]
    
    img = image.cpu().numpy().transpose((1, 2, 0))
    img = img * 0.5 + 0.5  # Un-normalize
    img = np.clip(img, 0, 1)
    axes[i].imshow(img, cmap="gray")
    axes[i].set_title(f"Original: {orig_label}\nPredicted: {pred_label}")
    axes[i].axis("off")
plt.show()

#Mobilenet


from torchvision.models import mobilenet_v2

# Define the pre-trained MobileNetV2 model
model = mobilenet_v2(pretrained=True)
num_ftrs = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_ftrs, 7)

# Move the model to the GPU if available
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())



# Train the model
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(train_loader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}")

# Save the model
torch.save(model.state_dict(), 'emotion_detection_model_wideresnet.pth')

model = mobilenet_v2(pretrained=True)
num_ftrs = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_ftrs, 7)
model=model.to(device)


# Load the pre-trained weights from a file
weights_path = "/content/emotion_detection_model_wideresnet.pth"
model.load_state_dict(torch.load(weights_path))
model=model.to(device)


# Test the model
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy: {100 * correct / total:.2f}%")


# Plot confusion matrix
from sklearn.metrics import confusion_matrix
import seaborn as sns

y_true = []
y_pred = []
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        y_true.extend(labels.cpu().numpy())
        y_pred.extend(predicted.cpu().numpy())

cm = confusion_matrix(y_true, y_pred)
class_labels = ['Angry', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad', 'Surprise']
sns.heatmap(cm, annot=True, xticklabels=class_labels, yticklabels=class_labels)
plt.show()

#Squeezenet

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.models import squeezenet1_0

# Define the pre-trained SqueezeNet model
model = squeezenet1_0(pretrained=True)
num_ftrs = model.classifier[1].in_channels
model.classifier[1] = nn.Conv2d(num_ftrs, 7, kernel_size=(1,1), stride=(1,1))
# Move the model to the GPU if available
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())

# Train the model
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(train_loader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}")

# Save the model
torch.save(model.state_dict(), 'emotion_detection_model_squeezenet.pth')

model = squeezenet1_0(pretrained=True)
num_ftrs = model.classifier[1].in_channels
model.classifier[1] = nn.Conv2d(num_ftrs, 7, kernel_size=(1,1), stride=(1,1))
model=model.to(device)


# Load the pre-trained weights from a file
weights_path = "/content/emotion_detection_model_squeezenet.pth"
model.load_state_dict(torch.load(weights_path))
model=model.to(device)


# Test the model
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy: {100 * correct / total:.2f}%")


# Plot confusion matrix
from sklearn.metrics import confusion_matrix
import seaborn as sns

y_true = []
y_pred = []
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        y_true.extend(labels.cpu().numpy())
        y_pred.extend(predicted.cpu().numpy())

cm = confusion_matrix(y_true, y_pred)
class_labels = ['Angry', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad', 'Surprise']
sns.heatmap(cm, annot=True, xticklabels=class_labels, yticklabels=class_labels)
plt.show()

#Shufflenet

# Define the pre-trained ShuffleNetV2 model
from torchvision.models import shufflenet_v2_x1_0
model = shufflenet_v2_x1_0(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 7)

# Move the model to the GPU if available
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())


# Train the model
num_epochs=50
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(train_loader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss / len(train_loader):.4f}")

# Save the model
torch.save(model.state_dict(), 'emotion_detection_model_shufflenet.pth')

# Test the model
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy: {100 * correct / total:.2f}%")


# Plot confusion matrix
from sklearn.metrics import confusion_matrix
import seaborn as sns

y_true = []
y_pred = []
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        y_true.extend(labels.cpu().numpy())
        y_pred.extend(predicted.cpu().numpy())

cm = confusion_matrix(y_true, y_pred)
class_labels = ['Angry', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad', 'Surprise']
sns.heatmap(cm, annot=True, xticklabels=class_labels, yticklabels=class_labels)
plt.show()

#Tensorboard

import torch
import torch.nn as nn
import torch.optim as optim
from torch.profiler import profile, record_function, ProfilerActivity
from torchvision.models import mobilenet_v2
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from google.colab import drive
drive.mount('/content/drive')

# Load the dataset from the zip file
import zipfile
with zipfile.ZipFile('/content/drive/MyDrive/archive.zip', 'r') as zip_ref:
    zip_ref.extractall('/content/dataset')

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, ToTensor, Normalize, RandomRotation, RandomHorizontalFlip
from torchvision.datasets import ImageFolder

from matplotlib import pyplot as plt
import numpy as np
import os
import random


# Parameters
IMG_HEIGHT = 48
IMG_WIDTH = 48
batch_size = 32
epochs = 50
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Directories
import torchvision
train_data_dir = '/content/dataset/train'
validation_data_dir = '/content/dataset/test'

mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]

# Define the data transforms for train and test sets
data_transforms = {
    'train': torchvision.transforms.Compose([
        torchvision.transforms.Resize(256),  # Resize the image to 256x256 pixels
        torchvision.transforms.RandomCrop(224),  # Crop a random 224x224 patch from the image
        torchvision.transforms.RandomRotation(30),  # Rotate the image randomly by up to 30 degrees
        torchvision.transforms.RandomHorizontalFlip(),  # Flip the image horizontally with a probability of 0.5
        torchvision.transforms.ToTensor(),  # Convert the image to a PyTorch tensor
        torchvision.transforms.Normalize(mean, std)  # Normalize the image using the mean and standard deviation
    ]),
    'test': torchvision.transforms.Compose([
        torchvision.transforms.Resize(256),  # Resize the image to 256x256 pixels
        torchvision.transforms.CenterCrop(224),  # Crop the center 224x224 patch from the image
        torchvision.transforms.ToTensor(),  # Convert the image to a PyTorch tensor
        torchvision.transforms.Normalize(mean, std)  # Normalize the image using the mean and standard deviation
    ])
}

# Datasets
train_dataset = ImageFolder(train_data_dir, data_transforms['train'])
test_dataset = ImageFolder(validation_data_dir, data_transforms['test'])

# DataLoaders
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

!pip install tensorboard
!pip install torch-tb-profiler

def profile_model(model, weights_path, log_dir):
    model.load_state_dict(torch.load(weights_path))
    model = model.to(device)

    writer = SummaryWriter(log_dir=log_dir)

    with torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=0, warmup=2, active=6, repeat=1),
        on_trace_ready=torch.profiler.tensorboard_trace_handler(log_dir),  # Pass log_dir instead of writer
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
        with_flops=True,
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    ) as prof:
        for i, (inputs, labels) in enumerate(test_loader):
            inputs, labels = inputs.to(device), labels.to(device)

            with record_function("forward"):
                outputs = model(inputs)
                loss = criterion(outputs, labels)

            if i >= 10:
                break

            prof.step()

    writer.close()


from torchvision.models import mobilenet_v2, shufflenet_v2_x1_0, resnet18, squeezenet1_0

resnet = resnet18(pretrained=True)
num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, 7)

mobilenet = mobilenet_v2(pretrained=True)
num_ftrs = mobilenet.classifier[1].in_features
mobilenet.classifier[1] = nn.Linear(num_ftrs, 7)

squeezenet = squeezenet1_0(pretrained=True)
num_ftrs = squeezenet.classifier[1].in_channels
squeezenet.classifier[1] = nn.Conv2d(num_ftrs, 7, kernel_size=(1,1), stride=(1,1))

shufflenet = shufflenet_v2_x1_0(pretrained=True)
num_ftrs = shufflenet.fc.in_features
shufflenet.fc = nn.Linear(num_ftrs, 7)

criterion = nn.CrossEntropyLoss()

models_and_paths = [
    (mobilenet, "/content/drive/MyDrive/emotion_detection_model_mobilenet.pth", "runs/profiler_mobilenet"),
    (shufflenet, "/content/drive/MyDrive/emotion_detection_model_shufflenet.pth", "runs/profiler_shufflenet"),
    (resnet, "/content/drive/MyDrive/emotion_detection_model_resnet.pth", "runs/profiler_resnet"),
    (squeezenet, "/content/drive/MyDrive/emotion_detection_model_squeezenet.pth", "runs/profiler_squeezenet"),
]

for model, weights_path, log_dir in models_and_paths:
    profile_model(model, weights_path, log_dir)


%load_ext tensorboard
%tensorboard --logdir runs