Simple Image Classification In Pytorch

1961 views May 18, 2024

ကျွန်တော် ဒီနေ့မှာတော့ Simple Image Classification ကို Pytorch ကို သုံးပြီး ဘယ်လို လုပ်ရလဲကို ပြောပြပေးသွားပါမယ်။

မှတ်ချက်။ ။ဒီ tutorial သည် beginners များအတွက် မဟုတ်ပါ။Neural Networks တွေနဲ့နည်းနည်းရင်းနှိးထားမှ ဖတ်လို့အဆင်ပြေပါလိမ့်မယ်။ဒါ့အပြင် စာကို သေချာဖတ်ပါ။ကျော်မဖတ်ပါနဲ့ ဘာလုပ်တယ်ဆိုတာကို အဲ့တာနားမလည်စေပါဘူး။ကျွန်တော်သည် Convolution ကဘာလဲ စတာတွေကို ပြောသွားမည်မဟုတ်ပါ။ကျွန်တော်သည် dataset ဘယ်လိုဆောက်မယ်နဲ့ pytorch ကိုသုံးပြီး ဘယ်လို train မယ်ကိုပဲ ပြပေးသွားမှာဖြစ်ပါတယ်။ဒီလို ရေးပေးရခြင်းသည် ကျွန်တော်တွေ့သလောက် pytorch tutorial ဆိုတာ မြန်မာလိုရေးထားတာရှားပါတယ်။Tensorflow,keras စတာတွေတော့ တွေ့ရပါတယ်။ဒါပေမဲ့ detailprojects blog ကို လာဖတ်တဲ့သူများဖြစ်ဖြစ် members ဝင်ထားသော သူငယ်ချင်းများဖြစ်ဖြစ် ပြောချင်တာက သူများ code တွေ content တွေကို မှီငြမ်းလို့ရပါတယ် ယူသုံးလို့လဲရပါတယ်။မှတ်ထားပေးရမွာက ကိုယ်ယူသုံးတဲ့ သူရဲ့ code ဖြစ်ဖြစ် tutorial content ဖြစ်ဖြစ် credit ဒါမှမဟုတ် တနေရာရာမှာ ဖော်ပြပေးပါ။ဒါလေးမှတ်သွားစေချင်ပါတယ်။

ဒီ tutorial တွင် ကျွန်တော်တို့ အရင်ဆုံး project structure ကို ပြောပါမယ်။အောက်မှာ ကြည်ေ့ပးပါ။

Project Folder

    data
        animals
            cats,dogs,pandas
    src
        models
        config.py
        dataset.py
        engine.py
        model.py
        test.py
        train.py

အဲ့တော့ ကျွန်တော်တို့သည် အရင်ဆုံး project folder ဆောက်လိုက်ပါ။အဲ့ထဲမှာ data နဲ့ src ဆိုတာရှိပါမယ်။data သည် images တွေထားဖို့ပါ။Animals ဆိုတဲ့ folder ထဲမှာ cats,dogs,pandas စသဖြင့် folders 3 ခုရှိပါတယ်။အဲ့ထဲမှာတော့ သက်ဆိုင်ရာပုံတွေရှိပါတယ်။
အဲ့တော့ မိတ်ဆွေက data folder ဆောက်ပြီးပြီဆိုရင် ဒီအောက်ကလင့်ခ်ကနေ download ဆွဲလိုက်ပြီး animals ထဲက folder 3 ခုကို အပေါ်က project structure တိုင်း ထားပေးပါ။

Click here to download data images           

ပြီးရင်တော့ src တွက် လိုအပ်တဲ့ code တွေကို အောက်က link ကနေ clone or download လုပ်ပြီး ယူပေးပါ။
Click here to download src files

ကဲ အကုန် လိုအပ်တာတွေ လုပ်ပြီးပြီဆိုတော့ အရင်ဆုံး src folder ထဲမှာ models folder ကို ဆောက်လိုက်ပါ။

models folder သည် ကျွန်တော်တို့ network ရဲ့ weight files တွေကိုသိမ်းဖို့ပါ။
အဲ့တော့ ကျွန်တော်သည် ဒီ code structure ဒီရက်ပိုင်း သဘောကျနေတာနဲ့ ပြီးတော့ မိတ်ဆွေတို့နဲ့ ပိုသင့်တော်ထင်လို့ တင်ပေးခြင်းဖြစ်ပါတယ်။ကျွန်တော်ရေးတာသည် သူ့ထက် နည်းနည်း complex ပိုဖြစ်ပါတယ်။ဒီ code style သည် youtube က data scientist တယောက်ရဲ့ code style ဖြစ်ပါတယ်။ကျွန်တော်ကတော့ လိုအပ်တာတွေနဲ့ ဒီ tutorial တွက် ကျွန်တော်ကြိုက်တဲ့ ပုံစံကို ပြန်ပြင်ထားပါတယ်။

အဲ့တော့ ကျွန်တော်တို့ configuration ဖြစ်တဲ့ config.py က စပြီး ကြည့်ရအောင်

Config.py

DATA_DIR = "../data/animals"
BATCH_SIZE = 128
IMAGE_WIDTH = 32
IMAGE_HEIGHT = 32
EPOCHS = 40
DEVICE = "cuda"

အဲ့တော့ ကျွန်တော်တို့ configuration ထဲမှာ DATA_DIR နဲ့ လိုအပ်တဲ့ batch_size တွေ ဘာတွေကို သတ်မှတ်လိုက်ပါတယ်။ကျွန်တော်ကတော့ GPU သုံးတာမို့လို့ cuda လို့ထားထားပါတယ်။မိတ်ဆွေတို့ cpu ပဲသုံးမယ်ဆို ဖြုတ်လိုက်ပါ။ပြီးရင်တော့ dataset.py ကနေပြီး dataset ကို စမ်းပြီး load လုပ်ကာ label များကို ယူရပါတယ်။

Dataset.py

from torchvision import transforms,datasets
from PIL import Image
import torch
import config
import glob
import os

from sklearn import preprocessing
import numpy as np 


class ClassificationDataset:
    def __init__(self,image_paths,targets,transform=None):
        self.image_paths = image_paths
        self.targets = targets
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self,item):
        image_path = self.image_paths[item]
        target = self.targets[item]
        image = Image.open(image_path).convert('RGB') # PIL image
        
        if self.transform is not None:
            image = self.transform(image)
        return {
            "images": image,
            "targets": torch.tensor(target,dtype=torch.long),
        }


if __name__ == "__main__":
    
    transform = transforms.Compose([
            transforms.Resize(size=(32,32)),
            transforms.ToTensor(),
            transforms.Normalize((0.45820624,0.43722707,0.39191988),(0.23130463,0.22692703,0.22379072))
    ])

    label_encoder = preprocessing.LabelEncoder()

    image_paths = glob.glob(os.path.join(config.DATA_DIR,"**/*.*"),recursive=True)
    targets = [x.split("/")[-2] for x in image_paths]
    label_encoded = label_encoder.fit_transform(targets)
    dataset = ClassificationDataset(image_paths,label_encoded,transform)
    print(np.unique(label_encoded))
    print(dataset[0]['images'].size())  
    print(dataset[0]['targets'])
    print(label_encoder.inverse_transform([dataset[0]['targets'].numpy()])[0]) 

ပထမဆုံး ကျွန်တော်တို့သည် လိုအပ်တဲ့ Dataset Class ဖြစ်တဲ့ ClassificationDataset ကို ဆောက်လိုက်ပြီး __getitem__ ထဲတွင် PIL နဲ့ image ဖတ်လိုက်ပြီး pytorch transforms ကနေ လိုအပ်တဲ့ tensor format နဲ့ image ကို normalize လုပ်လိုက်ပါတယ်။ပြီးရင်တော့ dataset ကနေ load လုပ်တိုင်း image နဲ့ target labels ကို ပြန်ပို့ပေးပါတယ်။အောက်မှာ dataset ကို စမ်းသပ်ထားတဲ့ code ကိုတွေ့နိုင်ပါတယ်။

နောက်တဆင့်သည် training function နဲ့ evaluation function များကို သတ်မှတ်ပြီး batch တခုချင်းဆီမှာ ဘယ်လောက် loss,accuracy ရလဲ တွက်ဖို့ပါ။

Engine.py

from tqdm import tqdm
import torch 
import config 


def metrics_batch(output, target):
    # get output class
    pred = output.argmax(dim=1, keepdim=True)
    
    # compare output class with target class
    corrects=pred.eq(target.view_as(pred)).sum().item()
    return corrects

def train_fn(model,data_loader,optimizer):
    model.train()
    fin_loss = 0.0
    tk = tqdm(data_loader,total=len(data_loader))
    for data in tk:
        for k,v in data.items():
            data[k] = v.to(config.DEVICE)
        
        optimizer.zero_grad()
        _,loss = model(**data)
        loss.backward()
        optimizer.step()
        fin_loss += loss.item()
    
    return fin_loss / len(data_loader)


def eval_fn(model,data_loader):
    model.eval()
    fin_loss = 0.0
    fin_accuracy = 0.0
    tk = tqdm(data_loader,total=len(data_loader))
    with torch.no_grad():
        for data in tk:
            for k,v in data.items():
                data[k] = v.to(config.DEVICE)
            
            batch_preds,loss = model(**data)
            fin_loss += loss.item()
            fin_accuracy += metrics_batch(batch_preds,data['targets'])
        
        return fin_accuracy / len(data_loader),fin_loss /len(data_loader) 

အဲ့တော့ အပေါ်က code တွေသည် training batch တွက် loss တွေကို တွက်ပြီး optimizer update လုပ်ပါတယ်။eval fn သည် validation batch အတွက် validation accuracy နဲ့ validation loss ကို တွက်ပြီး ပြန်ပို့ပေးပါတယ်။ဒီမှာ loss တွေသည် model ကနေ ပြန်ပို့ပေးထားတာတွေ့မှာပါ။အဲ့တော့ ကျွန်တော်တို့ model.py ကို ကြည့်ရအောင်

Model.py

import torch.nn as nn 
import torch.nn.functional as F 
import numpy as np 
import torch 

class SmallNet(nn.Module):
    def __init__(self,num_classes=3):
        super(SmallNet,self).__init__()
        # define the conv layers
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(2),
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.MaxPool2d(2),
        )
        
        self.num_flatten = 64 * 8 * 8
        self.fc1 = nn.Linear(self.num_flatten,512)
        self.out = nn.Linear(512,num_classes)
        
    
    def forward(self,images,targets=None):
        x = self.conv1(images)
        x = F.dropout(x, p=0.25)
        x = self.conv2(x)
        x = F.dropout(x, p=0.25)
        x = x.view(-1,self.num_flatten)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, p=0.25)
        x = self.out(x)
        if targets is not None:
            loss = nn.CrossEntropyLoss()
            loss = loss(x, targets)
            return x,loss
        return x,None 

if __name__ == "__main__":
    model = SmallNet(num_classes=3)
    img = torch.rand((1,3,32,32))
    x,loss = model(img,torch.tensor([1],dtype=torch.long))

Model.py သည် small vgg architecture သဘောမျိုးသာရေးထားတာတွေ့ရမှာပါ။ဒီမှာ pytorch neural network ကို မရှင်းတော့ဘူးနော်။Intermediate reader ဆိုရင်တော့ ဒီ code ကိုကြည့်တာနဲ့ နည်းနည်းတော့ သဘောပေါက်မှာပါ။

အဲ့တော့ ကျွန်တော်တို့ လိုအပ်တဲ့ dataset class တွေ config တွေ neural network model တွေပြင်ပြီးသွားပြီမို့ စ train မှာမို့ train.py ကိုကြည့်ရအောင်

Train.py


import os
import glob
import torch
import numpy as np

from torchvision import transforms,datasets

from sklearn import preprocessing
from sklearn import model_selection

import config
import dataset
import engine
from model import SmallNet


def run_training():
    transform = transforms.Compose([
            transforms.Resize(size=(32,32)),
            transforms.ToTensor(),
            transforms.Normalize((0.45820624,0.43722707,0.39191988),(0.23130463,0.22692703,0.22379072))
    ])

    label_encoder = preprocessing.LabelEncoder()

    image_paths = glob.glob(os.path.join(config.DATA_DIR,"**/*.*"),recursive=True)
    targets = [x.split("/")[-2] for x in image_paths]
    label_encoded = np.array(label_encoder.fit_transform(targets))
    
    (train_images,test_images,train_labels,test_labels) = model_selection.train_test_split(image_paths,label_encoded,test_size=0.2,random_state=0)
    # print(len(train_images))
    # print(len(train_labels))

    train_dataset = dataset.ClassificationDataset(train_images,train_labels,transform)
    train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=config.BATCH_SIZE,shuffle=True)
    
    test_dataset = dataset.ClassificationDataset(test_images,test_labels,transform)
    test_loader = torch.utils.data.DataLoader(test_dataset,batch_size=config.BATCH_SIZE,shuffle=False)

    model =SmallNet(num_classes=3)
    model.to(config.DEVICE)

    opt = torch.optim.Adam(model.parameters(),lr=3e-4)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        opt, factor=0.8, patience=5, verbose=True
    )

    for epoch in range(config.EPOCHS):
        
        train_loss = engine.train_fn(model,train_loader,opt)
        val_accuracy,val_loss = engine.eval_fn(model,test_loader)

        print(
            f"Epoch={epoch}, Train Loss={train_loss}, Test Loss={val_loss} Accuracy={val_accuracy}"
        )
        
        scheduler.step(val_loss)

    print("Saved model...")
    torch.save(model.state_dict(),"./models/weights_latest.pt")    

if __name__ == "__main__":
    run_training()

အဲ့တော့ ကျွန်တော်တို့သည် run_training function ထဲမှာ dataset ကို load လုပ်ပြီး pytorch Dataloader ကနေတဆင့် engine.py မှာရေးထားတဲ့ training function တွေ evaluation function တွေ ပြန်ပို့ပေးကာ loss များ ရယူပြီး train လိုက်ပါတယ်။ကျွန်တော်ဒီမှာ နောက်ဆုံး weights ကို save ထားပါတယ်။လုပ်တော့မလုပ်သင့်ပေမယ့် tutorial purpose ပဲမို့ ဒီလိုလေးပဲထားခဲ့ပါတော့မယ်။

အဲ့တော့ ကျွန်တော်တို့ train.py ကို run လိုက်မယ်ဆို အောက်ပါအတိုင်း တွေ့ရပါမယ်။

ဒီမှာ ကျွန်တော်တို့ epoch 40 ထိtrian ပြီးတဲ့ခါ accuracy လေးက ကောင်းပေမယ့် train loss ,validation loss သည် ကွာနေတာ တွေ့ရပါလိမ့်မယ်။ဒါသည် overfitting ဖြစ်နေပြီး ကောင်းတဲ့အချက်တော့မဟုတ်ပါ။training data အရမ်းနည်းလို့ဖြစ်ပါတယ်။အဲ့တော့ ကိုယ့်ဘာ့သာကိုယ် ဘယ်လိုကောင်းအောင် လုပ်မလဲ စမ်းကြည့်ကြပါ။ပြီးရင်တော့ ကျွန်တော်တို့ model deployment test.py ကိုကြည့်ရအောင်။

Test.py

import numpy as np
from model import SmallNet
from PIL import Image 
import matplotlib.pyplot as plt 
from torchvision import transforms
import torch 
import argparse

ap = argparse.ArgumentParser()
ap.add_argument("-i","--input",required=True,help="Path to input image")
args = vars(ap.parse_args())

transform = transforms.Compose([
            transforms.Resize(size=(32,32)),
            transforms.ToTensor(),
            transforms.Normalize((0.45820624,0.43722707,0.39191988),(0.23130463,0.22692703,0.22379072))
    ])

labels = ['cat','dog','panda']
path2modelweight = "./models/weights_latest.pt"

model = SmallNet(num_classes=3)
model.load_state_dict(torch.load(path2modelweight))

model.eval()
image = Image.open(args['input'])

with torch.no_grad():
    x = transform(image)
    y,_ = model(x.unsqueeze(0))
    y_out = torch.softmax(y,dim=1).numpy()
    y_out = np.argmax(y_out)
    print("Label->",labels[y_out])
    fig,ax = plt.subplots()
    ax.imshow(image)
    plt.show()


ကျွန်တော်တို့ test.py တွင် argument ကဝင်လာတဲ့ image file ကို read လိုက်ပြီး model ထဲထည့်တာ prediction ကိုယူတာလေးဖြစ်ပါတယ်။အားလုံးနားလည်မယ်လို့မျှော်လင့်ပါတယ်။

အဲ့တော့ train ပြီးလို့ run မည်ဆိုလျှင် အောက်ပါတိုင်း run လို့ရပါတယ်။


ဒီမှာ jupyter notebook မသုံးပဲ ဒါမျိုးရေးရခြင်းသည် production တွက်ကော ကေင်းပါတယ်။Python code သည် python code ပါပဲ ဘယ်မှာရေးရေးတူတူပါပဲ။အားလုံးကျေးဇူးတင်ပါတယ်။