model_development.py

# -*- coding: utf-8 -*-
"""Model_Development

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1BPWPi-oYa82w42fjYxMY7SVOU2XwedyW
"""

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.utils import image_dataset_from_directory
import os
import numpy as np
import keras
from keras import layers
from tensorflow import data as tf_data
import matplotlib.pyplot as plt

!unzip "/content/drive/MyDrive/Caitlin Bodzy/Data/constellations.zip" -d "/content"



# Preparing our data
directory = "/content/content/constellations2"
image_size = (128, 128)
batch_size = 64

train_ds, val_ds = image_dataset_from_directory(
    directory,
    validation_split=0.2,
    subset="both",
    seed=1337,
    image_size=image_size,
    batch_size=batch_size,

)

plt.figure(figsize=(10, 10))
for images, labels in train_ds.take(1):
    for i in range(9):
        ax = plt.subplot(3, 3, i + 1)
        plt.imshow(np.array(images[i]).astype("uint8"))
        plt.title(int(labels[i]))
        plt.axis("off")

def make_model(input_shape, num_classes):
    inputs = keras.Input(shape=input_shape)

    # Entry block
    x = layers.Rescaling(1.0 / 255)(inputs)
    x = layers.Conv2D(128, 3, strides=2, padding="same")(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)

    previous_block_activation = x  # Set aside residual

    for size in [256, 512, 728]:
        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(size, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(size, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.MaxPooling2D(3, strides=2, padding="same")(x)

        # Project residual
        residual = layers.Conv2D(size, 1, strides=2, padding="same")(
            previous_block_activation
        )
        x = layers.add([x, residual])  # Add back residual
        previous_block_activation = x  # Set aside next residual

    x = layers.SeparableConv2D(1024, 3, padding="same")(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)

    x = layers.GlobalAveragePooling2D()(x)
    if num_classes == 2:
        units = 1
    else:
        units = num_classes

    x = layers.Dropout(0.25)(x)
    # We specify activation=None so as to return logits
    outputs = layers.Dense(units, activation=None)(x)
    return keras.Model(inputs, outputs)

num_classes = 89
model = make_model(input_shape=image_size + (3,), num_classes=num_classes)
keras.utils.plot_model(model, show_shapes=True)

epochs = 25

callbacks = [
    keras.callbacks.ModelCheckpoint("save_at_{epoch}.keras"),
]
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)
model.fit(
    train_ds,
    epochs=epochs,
    callbacks=callbacks,
    validation_data=val_ds,

)

"""#Pytorch time"""

import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from torchvision.datasets import ImageFolder
from torchvision import datasets, models, transforms
from tqdm import tqdm
import torch.optim as optim
import pickle

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# can use ImageFolderDataset
path = "/content/content/constellations2"
dataset = ImageFolder(path, transform = transform)

train_dataset, test_dataset = torch.utils.data.random_split(dataset, [0.8, 0.2])

batch_size = 64
trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,
                                          shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size,
                                         shuffle=False, num_workers=2)

def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=10, device="cuda"):
    """
    Trains a PyTorch model and prints training + validation loss/accuracy each epoch.

    Args:
        model: nn.Module - your model
        train_loader: DataLoader - training data
        val_loader: DataLoader - validation data
        criterion: loss function (e.g. nn.CrossEntropyLoss)
        optimizer: optimizer (e.g. Adam, SGD)
        num_epochs: int - number of training epochs
        device: 'cuda' or 'cpu'
    """

    model.to(device)
    # Initialize history dictionary
    history = {
        "train_loss": [],
        "train_acc": [],
        "val_loss": [],
        "val_acc": []
    }
    best_accuracy = 0
    for epoch in range(num_epochs):
        # --- Training phase ---
        model.train()  # set model to training mode
        running_loss = 0.0
        correct = 0
        total = 0
        print(f"\nEpoch [{epoch+1}/{num_epochs}]")
        train_pbar = tqdm(train_loader, desc="Training", leave=False)
        for inputs, labels in train_pbar:
            # Move data to GPU/CPU
            inputs, labels = inputs.to(device), labels.to(device)

            # Zero gradients for this batch
            optimizer.zero_grad()

            # Forward pass
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            # Backward pass + optimization
            loss.backward()
            optimizer.step()

            # Update training loss
            running_loss += loss.item() * inputs.size(0)

            # Get predictions (highest logit = predicted class)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            # live update in progress bar
            train_pbar.set_postfix(loss=loss.item())
        epoch_train_loss = running_loss / len(train_loader.dataset)
        epoch_train_acc = 100 * correct / total
        history["train_loss"].append(epoch_train_loss)
        history["train_acc"].append(epoch_train_acc)
        # --- Validation phase ---
        model.eval()
        val_loss, val_correct, val_total = 0.0, 0, 0

        val_pbar = tqdm(val_loader, desc="Validating", leave=False)
        with torch.no_grad():
            for inputs, labels in val_pbar:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                val_loss += loss.item() * inputs.size(0)
                _, predicted = torch.max(outputs, 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()

                # live update in progress bar
                val_pbar.set_postfix(loss=loss.item())

        epoch_val_loss = val_loss / len(val_loader.dataset)
        epoch_val_acc = 100 * val_correct / val_total
        if epoch_val_acc > best_accuracy:
          best_accuracy = epoch_val_acc
          torch.save(model.state_dict(), "best_model_params.pt")
        history["val_loss"].append(epoch_val_loss)
        history["val_acc"].append(epoch_val_acc)
        # --- Print results at end of epoch ---
        print(f"Epoch [{epoch+1}/{num_epochs}] "
              f"Train Loss: {epoch_train_loss:.4f}, Train Acc: {epoch_train_acc:.2f}% "
              f"| Val Loss: {epoch_val_loss:.4f}, Val Acc: {epoch_val_acc:.2f}%")
    print("\n✅ Training complete.")
    return model, history

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 5, stride=2, padding=2)  # downsample early
        self.conv2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, stride=2, padding=1)
        self.conv4 = nn.Conv2d(128, 256, 3, stride=2, padding=1)

        # Global average pooling (output shape = batch × 256)
        self.gap = nn.AdaptiveAvgPool2d((1, 1))

        # Small fully connected head
        self.fc1 = nn.Linear(256, 128)
        self.fc2 = nn.Linear(128, 89)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = self.gap(x)              # -> (batch, 256, 1, 1)
        x = torch.flatten(x, 1)      # -> (batch, 256)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)
model = Net().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
num_epochs = 50
model, history = train_model(model, trainloader, testloader, criterion, optimizer, num_epochs=num_epochs, device=device)

with open('custom_cnn_history.pkl', 'wb') as f:
    pickle.dump(history, f)

"""Results for 50 epochs:

Using device: cuda

Epoch [1/50]
Epoch [1/50] Train Loss: 4.3550, Train Acc: 1.81% | Val Loss: 4.0249, Val Acc: 2.54%

Epoch [2/50]
Epoch [2/50] Train Loss: 3.8108, Train Acc: 3.31% | Val Loss: 3.6244, Val Acc: 3.45%

Epoch [3/50]
Epoch [3/50] Train Loss: 3.4859, Train Acc: 6.98% | Val Loss: 3.3346, Val Acc: 5.63%

Epoch [4/50]
Epoch [4/50] Train Loss: 3.2355, Train Acc: 8.43% | Val Loss: 3.1397, Val Acc: 12.89%

Epoch [5/50]
Epoch [5/50] Train Loss: 3.0890, Train Acc: 11.51% | Val Loss: 3.0306, Val Acc: 12.16%

Epoch [6/50]
Epoch [6/50] Train Loss: 3.1845, Train Acc: 9.11% | Val Loss: 3.0034, Val Acc: 19.60%

Epoch [7/50]
Epoch [7/50] Train Loss: 2.8859, Train Acc: 15.45% | Val Loss: 2.6902, Val Acc: 15.97%

Epoch [8/50]
Epoch [8/50] Train Loss: 2.7958, Train Acc: 16.54% | Val Loss: 2.7449, Val Acc: 16.70%

Epoch [9/50]
Epoch [9/50] Train Loss: 2.5519, Train Acc: 22.34% | Val Loss: 2.7096, Val Acc: 17.79%

Epoch [10/50]
Epoch [10/50] Train Loss: 2.4776, Train Acc: 23.33% | Val Loss: 2.4054, Val Acc: 24.68%

Epoch [11/50]
Epoch [11/50] Train Loss: 2.2965, Train Acc: 29.50% | Val Loss: 2.2966, Val Acc: 23.96%

Epoch [12/50]
Epoch [12/50] Train Loss: 2.0906, Train Acc: 33.57% | Val Loss: 2.0627, Val Acc: 36.30%

Epoch [13/50]
Epoch [13/50] Train Loss: 2.0251, Train Acc: 37.52% | Val Loss: 1.9627, Val Acc: 37.57%

Epoch [14/50]
Epoch [14/50] Train Loss: 1.7805, Train Acc: 44.09% | Val Loss: 1.8681, Val Acc: 40.83%

Epoch [15/50]
Epoch [15/50] Train Loss: 1.7417, Train Acc: 43.04% | Val Loss: 1.5517, Val Acc: 51.36%

Epoch [16/50]
Epoch [16/50] Train Loss: 1.5492, Train Acc: 47.80% | Val Loss: 1.5065, Val Acc: 51.18%

Epoch [17/50]
Epoch [17/50] Train Loss: 1.5765, Train Acc: 48.39% | Val Loss: 1.4849, Val Acc: 52.99%

Epoch [18/50]
Epoch [18/50] Train Loss: 1.3129, Train Acc: 57.68% | Val Loss: 1.1928, Val Acc: 58.62%

Epoch [19/50]
Epoch [19/50] Train Loss: 1.6431, Train Acc: 46.35% | Val Loss: 1.3933, Val Acc: 54.81%

Epoch [20/50]
Epoch [20/50] Train Loss: 1.1105, Train Acc: 64.20% | Val Loss: 1.2303, Val Acc: 57.89%

Epoch [21/50]
Epoch [21/50] Train Loss: 1.1097, Train Acc: 61.53% | Val Loss: 0.9919, Val Acc: 65.88%

Epoch [22/50]
Epoch [22/50] Train Loss: 0.9873, Train Acc: 66.65% | Val Loss: 0.9837, Val Acc: 65.52%

Epoch [23/50]
Epoch [23/50] Train Loss: 0.8362, Train Acc: 70.00% | Val Loss: 1.1519, Val Acc: 62.61%

Epoch [24/50]
Epoch [24/50] Train Loss: 0.9181, Train Acc: 67.47% | Val Loss: 0.9100, Val Acc: 70.05%

Epoch [25/50]
Epoch [25/50] Train Loss: 1.0156, Train Acc: 65.52% | Val Loss: 0.9139, Val Acc: 67.70%

Epoch [26/50]
Epoch [26/50] Train Loss: 0.7847, Train Acc: 73.72% | Val Loss: 0.6174, Val Acc: 82.94%

Epoch [27/50]
Epoch [27/50] Train Loss: 0.8881, Train Acc: 70.41% | Val Loss: 1.0006, Val Acc: 60.80%

Epoch [28/50]
Epoch [28/50] Train Loss: 0.9419, Train Acc: 69.05% | Val Loss: 0.6462, Val Acc: 83.85%

Epoch [29/50]
Epoch [29/50] Train Loss: 0.6375, Train Acc: 78.84% | Val Loss: 0.6832, Val Acc: 76.23%

Epoch [30/50]
Epoch [30/50] Train Loss: 0.5708, Train Acc: 78.66% | Val Loss: 0.6045, Val Acc: 78.04%

Epoch [31/50]
Epoch [31/50] Train Loss: 0.5209, Train Acc: 81.15% | Val Loss: 0.8075, Val Acc: 76.77%

Epoch [32/50]
Epoch [32/50] Train Loss: 0.6596, Train Acc: 77.03% | Val Loss: 0.5404, Val Acc: 82.21%

Epoch [33/50]
Epoch [33/50] Train Loss: 0.4934, Train Acc: 84.05% | Val Loss: 0.3702, Val Acc: 87.66%

Epoch [34/50]
Epoch [34/50] Train Loss: 0.5477, Train Acc: 79.97% | Val Loss: 0.4291, Val Acc: 85.48%

Epoch [35/50]
Epoch [35/50] Train Loss: 0.4817, Train Acc: 83.55% | Val Loss: 0.6402, Val Acc: 76.23%

Epoch [36/50]
Epoch [36/50] Train Loss: 0.7414, Train Acc: 73.95% | Val Loss: 0.7658, Val Acc: 74.23%

Epoch [37/50]
Epoch [37/50] Train Loss: 0.4875, Train Acc: 83.64% | Val Loss: 0.5194, Val Acc: 79.13%

Epoch [38/50]
Epoch [38/50] Train Loss: 0.7001, Train Acc: 74.94% | Val Loss: 0.4669, Val Acc: 82.40%

Epoch [39/50]
Epoch [39/50] Train Loss: 0.4260, Train Acc: 86.27% | Val Loss: 0.4313, Val Acc: 84.21%

Epoch [40/50]
Epoch [40/50] Train Loss: 0.4358, Train Acc: 84.05% | Val Loss: 1.0168, Val Acc: 66.79%

Epoch [41/50]
Epoch [41/50] Train Loss: 0.4003, Train Acc: 85.50% | Val Loss: 0.4285, Val Acc: 85.12%

Epoch [42/50]
Epoch [42/50] Train Loss: 0.3837, Train Acc: 86.91% | Val Loss: 0.6751, Val Acc: 78.40%

Epoch [43/50]
Epoch [43/50] Train Loss: 0.4011, Train Acc: 85.46% | Val Loss: 1.9752, Val Acc: 55.54%

Epoch [44/50]
Epoch [44/50] Train Loss: 0.9230, Train Acc: 73.18% | Val Loss: 0.4821, Val Acc: 84.94%

Epoch [45/50]
Epoch [45/50] Train Loss: 0.3531, Train Acc: 89.13% | Val Loss: 0.3267, Val Acc: 90.38%

Epoch [46/50]
Epoch [46/50] Train Loss: 0.3086, Train Acc: 89.40% | Val Loss: 0.3978, Val Acc: 85.30%

Epoch [47/50]
Epoch [47/50] Train Loss: 0.3021, Train Acc: 89.49% | Val Loss: 0.3236, Val Acc: 88.02%

Epoch [48/50]
Epoch [48/50] Train Loss: 0.3974, Train Acc: 85.27% | Val Loss: 0.3593, Val Acc: 87.11%

Epoch [49/50]
Epoch [49/50] Train Loss: 0.2615, Train Acc: 90.39% | Val Loss: 0.2154, Val Acc: 92.38%

Epoch [50/50]
                                                                     Epoch [50/50] Train Loss: 0.2437, Train Acc: 92.03% | Val Loss: 0.3268, Val Acc: 88.38%

✅ Training complete.
"""

import matplotlib.pyplot as plt

# summarize history for accuracy
plt.plot(history['train_acc'])
plt.plot(history['val_acc'])
plt.title('custom cnn model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['Train', 'Validation'], loc='upper left')
plt.show()
# summarize history for loss
plt.plot(history['train_loss'])
plt.plot(history['val_loss'])
plt.title('custom cnn model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['Train', 'Validation'], loc='upper left')
plt.show()

"""#resnet"""

#Different model
model_ft = models.resnet18(weights='IMAGENET1K_V1')
num_ftrs = model_ft.fc.in_features
# Here the size of each output sample is set to 2.
# Alternatively, it can be generalized to ``nn.Linear(num_ftrs, len(class_names))``.
model_ft.fc = nn.Linear(num_ftrs, 89)

from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms
import torch

# Standard ResNet preprocessing
image_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])  # ImageNet stats
])
batch_size = 64
# can use ImageFolderDataset
path = "/content/content/constellations2"
dataset = ImageFolder(path, transform = image_transforms)
train_dataset, test_dataset = torch.utils.data.random_split(dataset, [0.8, 0.2])
trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,
                                          shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size,
                                         shuffle=False, num_workers=2)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)
model = Net().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
model, history = train_model(model_ft, trainloader, testloader, criterion, optimizer, num_epochs=50, device=device)

# summarize history for accuracy
plt.plot(history['train_acc'])
plt.plot(history['val_acc'])
plt.title('resnet model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['Train', 'Validation'], loc='upper left')
plt.show()
# summarize history for loss
plt.plot(history['train_loss'])
plt.plot(history['val_loss'])
plt.title('resnet model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['Train', 'Validation'], loc='upper left')
plt.show()

"""Chat gpt improved resnet18 model"""

# ==== Imports ====
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader, random_split
from tqdm import tqdm
import matplotlib.pyplot as plt

# ==== 1. Device setup ====
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# ==== 2. Data setup with augmentations ====
path = "/content/content/constellations2"

image_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])  # ImageNet normalization
])

dataset = ImageFolder(path, transform=image_transforms)

# Split dataset into train/test
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

# DataLoaders
batch_size = 64
trainloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
testloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)

# ==== 3. Model setup ====
model_ft = models.resnet18(weights='IMAGENET1K_V1')
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 89)  # 89 output classes
model_ft = model_ft.to(device)

# ==== 4. Loss and optimizer ====
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model_ft.parameters(), lr=1e-4, weight_decay=1e-4)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

# ==== 5. Training function ====
def train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device):
    history = {'train_loss': [], 'val_loss': [], 'train_acc': [], 'val_acc': []}

    for epoch in range(num_epochs):
        print(f"\nEpoch [{epoch+1}/{num_epochs}]")
        model.train()
        train_loss, correct, total = 0.0, 0, 0

        for inputs, labels in tqdm(train_loader, desc="Training", leave=False):
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        scheduler.step()

        train_acc = 100 * correct / total
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)

        history['train_loss'].append(train_loss / len(train_loader.dataset))
        history['val_loss'].append(val_loss)
        history['train_acc'].append(train_acc)
        history['val_acc'].append(val_acc)

        print(f"Train Loss: {train_loss/len(train_loader.dataset):.4f} | "
              f"Train Acc: {train_acc:.2f}% | "
              f"Val Loss: {val_loss:.4f} | "
              f"Val Acc: {val_acc:.2f}%")

    return model, history

# ==== 6. Validation function ====
def evaluate(model, loader, criterion, device):
    model.eval()
    loss_total, correct, total = 0.0, 0, 0
    with torch.no_grad():
        for inputs, labels in loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss_total += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    avg_loss = loss_total / len(loader.dataset)
    acc = 100 * correct / total
    return avg_loss, acc

# ==== 7. Train model ====
model_ft, history = train_model(
    model_ft,
    trainloader,
    testloader,
    criterion,
    optimizer,
    scheduler,
    num_epochs=30,
    device=device
)

# ==== 8. Visualization ====
plt.figure(figsize=(10,5))
plt.plot(history['train_acc'], label='Train Accuracy')
plt.plot(history['val_acc'], label='Validation Accuracy')
plt.title('ResNet18 Model Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.legend()
plt.show()

plt.figure(figsize=(10,5))
plt.plot(history['train_loss'], label='Train Loss')
plt.plot(history['val_loss'], label='Validation Loss')
plt.title('ResNet18 Model Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()