Image Segmentation - Single Layer Model

# models/single_conv.py
import torch.nn as nn

class SingleConvolution(nn.Module):
    def __init__(self):
        super(SingleConvolution, self).__init__()
        self.conv_layer1 = nn.Sequential (
            nn.Conv2d(3, 23, kernel_size = 3, padding=1),
            nn.ReLU()
        )
    def forward(self, x):
        output = self.conv_layer1(x)
        return output

# main.py
from basics import *
from datautils import * 
from metrics import * 
from evaluate import *
from train import * 
from models import * 

batch_size = 6
train_set, test_set = import_drone_dataset()
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
model = SingleConvolution().to(device)
print(summary(model, (3, 600, 400)))

# train.py
from basics import * 
from metrics import *

def train(
    model : nn.Module, 
    epochs : int,
    train_loader : DataLoader,
    loss_func, optimizer,
    acc_metric = pixel_accuracy
):
    torch.cuda.empty_cache()
    train_losses = [] 
    train_acc = []
    model.to(device)
    start = time.time()
    for _ in range(epochs):
        print(f"EPOCH {_+1} training begins...")
        train_start = time.time()
        model.train()
        train_accuracy = 0 
        train_loss = 0
        for i, data in enumerate(tqdm(train_loader)):
            img, mask = data 
            img = img.to(device)
            mask = mask.to(device)
            output = model(img)

            train_accuracy += acc_metric(output, mask)
            loss = loss_func(output, mask)
            train_loss += loss.item()
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        print(f"Train epoch {_+1} / {epochs}",
              f"Training Loss {train_loss/len(train_loader):.4f}",
              f"Training Accr {train_accuracy/len(train_loader):.4f}",
              f"Training Time {(time.time()-train_start)/60:.2f} min")
    history = {'train_loss' : train_losses, 'train_acc' :train_acc}
    print(f"Total training time {(time.time()-start)/60:.2f} minutes taken")
    return history

# main.py
train(
    model = model, 
    epochs = 5,
    train_loader = train_loader, 
    loss_func = nn.CrossEntropyLoss(), 
    optimizer = torch.optim.Adam(model.parameters(), lr=0.003)
)
evaluator = ModelEvaluation(model, test_set, pixel_accuracy)
evaluator.evaluate_all()