Skip to content

Latest commit

 

History

History
65 lines (54 loc) · 1.54 KB

README.md

File metadata and controls

65 lines (54 loc) · 1.54 KB
Raw

Pytorch_ImageClassification

  • Pytorch
  • Classification
  • AlexNet, VGGNet, ResNet

Version

  • torch : 1.8.0+cu111
  • torchvision : 0.9.0+cu111
  • python : 3.7.6
  • numpy : 1.18.2

Usage Exaple

1. Set Parameter

# parameter
num_classes = 2
num_epoch = 50
batch_size = 64
learning_rate = 0.001
image_size = 224

train_data_path = '../dataset/testset/train'
test_data_path = '../dataset/testset/test'
save_path = '../backup/new_model.pth'

transform = transforms.Compose([
    transforms.Resize((image_size, image_size)),
    transforms.ToTensor()
])
  • num_classes : number of class
  • num_epoch : train epoch
  • batch_size : train dataset batch size
  • image_size : check model input size

2. Select model

# load model
model = QResNet152(num_classes=num_classes)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
  • QAlexNet (image size : 227x227)
  • QVGGNet (224x224)
  • QResNet34, QResNet152 (224x224)

3. Training & Test

train_graph

4. Grad-CAM

# Grad-CAM
from GradCAM import GetGradCAMModule

gc_model = GetGradCAMModule(load_model, 'conv5_x')
gc_model.eval()
  • GetGradCAMModule(model, target_layer)
  • model : base CNN model
  • target_layer : layer to visualize with GradCAM

9 45