Image Classification with ResNet50 in Pytorch

https://blog.devgenius.io/resnet50-6b42934db431

ResNet50

 

1. set env

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import os

import torch
import torchvision
from torchvision import datasets, models, transforms
import torch.nn as nn
from torch.nn import functional as F
import torch.optim as optim

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   
os.environ["CUDA_VISIBLE_DEVICES"]="0"

2. data generator

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])

data_transforms = {
    'train':
    transforms.Compose([
        transforms.Resize((224,224)),
        transforms.RandomAffine(0, shear=10, scale=(0.8,1.2)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        normalize
    ]),
    'val':
    transforms.Compose([
        transforms.Resize((224,224)),
        transforms.ToTensor(),
        normalize
    ]),
    # 'test':
    # transforms.Compose([
    #     transforms.Resize((224,224)),
    #     transforms.ToTensor(),
    #     normalize
    # ]),
}

input_path = 'D:/0616dataset/split9'

image_datasets = {
    'train': 
    datasets.ImageFolder(input_path + '/train', data_transforms['train']),
    'val': 
    datasets.ImageFolder(input_path + '/val', data_transforms['val']),
    # 'test': 
    # datasets.ImageFolder(input_path + 'test', data_transforms['test'])
}

dataloaders = {
    'train':
    torch.utils.data.DataLoader(image_datasets['train'],
                                batch_size=32,
                                shuffle=True,
                                num_workers=0),  
    'val':
    torch.utils.data.DataLoader(image_datasets['val'],
                                batch_size=32,
                                shuffle=False,
                                num_workers=0),
    # 'test':
    # torch.utils.data.DataLoader(image_datasets['test'],
    #                             batch_size=32,
    #                             shuffle=False,
    #                             num_workers=0)
}

3. set gpu device

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device

4. load pretrained resnet50

model = models.resnet50(pretrained=True).to(device)
    
for param in model.parameters():
    param.requires_grad = False   
    
model.fc = nn.Sequential(
               nn.Linear(2048, 128),
               nn.ReLU(inplace=True),
               nn.Linear(128, 1253)).to(device)