You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
193 lines
4.9 KiB
193 lines
4.9 KiB
import torch
|
|
import numpy as np
|
|
import cv2, os, sys
|
|
import pandas as pd
|
|
from torch.utils.data import Dataset
|
|
from matplotlib import pyplot as plt
|
|
from torch.utils.data import ConcatDataset, DataLoader, Subset
|
|
import torch.nn as nn
|
|
import torchvision.transforms as transforms
|
|
from torchvision.datasets import DatasetFolder
|
|
from PIL import Image
|
|
import torchvision.models as models
|
|
batch_size = 32
|
|
num_epoch = 1
|
|
|
|
torch.cuda.set_device(1)
|
|
train_tfm = transforms.Compose([
|
|
transforms.Grayscale(),
|
|
transforms.RandomHorizontalFlip(),
|
|
transforms.RandomResizedCrop((68,68)),
|
|
transforms.ToTensor(),
|
|
#transforms.RandomResizedCrop((40,30)),
|
|
#transforms.TenCrop((40,30)),
|
|
#transforms.Normalize(0.5,0.5),
|
|
])
|
|
|
|
test_tfm = transforms.Compose([
|
|
transforms.Grayscale(),
|
|
transforms.ToTensor()
|
|
])
|
|
'''
|
|
class Classifier(nn.Module):
|
|
def __init__(self):
|
|
super(Classifier, self).__init__()
|
|
|
|
self.cnn_layers = nn.Sequential(
|
|
#input_size(1,30,40)
|
|
nn.Conv2d(1, 16, 3, 1), #output_size(16,28,38)
|
|
nn.BatchNorm2d(16),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.2),
|
|
nn.MaxPool2d(kernel_size = 2), #output_size(16,14,19)
|
|
|
|
nn.Conv2d(16, 24, 3, 1), #output_size(24,12,17)
|
|
nn.BatchNorm2d(24),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.2),
|
|
nn.MaxPool2d(kernel_size = 2), #output_size(24,6,8)
|
|
|
|
nn.Conv2d(24, 32, 3, 1), #output_size(32,4,6)
|
|
nn.BatchNorm2d(32),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.2),
|
|
nn.MaxPool2d(kernel_size = 2) #ouput_size(32,2,3)
|
|
|
|
|
|
)
|
|
self.fc_layers = nn.Sequential(
|
|
nn.Linear(32 * 2 * 3, 32),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.2),
|
|
|
|
|
|
nn.Linear(32,8)
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.cnn_layers(x)
|
|
|
|
x = x.flatten(1)
|
|
|
|
x = self.fc_layers(x)
|
|
|
|
return x
|
|
'''
|
|
def main():
|
|
train_set = DatasetFolder("pose_data2/train", loader=lambda x: Image.open(x), extensions="bmp", transform=train_tfm)
|
|
test_set = DatasetFolder("pose_data2/test", loader=lambda x: Image.open(x), extensions="bmp", transform=test_tfm)
|
|
valid_set = DatasetFolder("pose_data2/val", loader=lambda x: Image.open(x), extensions="bmp", transform=test_tfm)
|
|
|
|
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
|
|
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=True)
|
|
valid_loader = DataLoader(valid_set, batch_size=batch_size, shuffle=True)
|
|
|
|
model_path = "model.ckpt"
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model = models.resnet50()
|
|
model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3,
|
|
bias=False)
|
|
model.fc = nn.Linear(2048, 8)
|
|
model = model.to(device)
|
|
print(model)
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
|
|
criterion = nn.CrossEntropyLoss()
|
|
|
|
best_acc = -1
|
|
for epoch in range(num_epoch):
|
|
##Training
|
|
running_loss = 0.0
|
|
total = 0
|
|
correct = 0
|
|
for i, data in enumerate(train_loader):
|
|
inputs, labels = data
|
|
inputs = inputs.to(device)
|
|
labels = labels.to(device)
|
|
optimizer.zero_grad()
|
|
|
|
outputs = model(inputs)
|
|
|
|
loss = criterion(outputs, labels)
|
|
loss.backward()
|
|
|
|
optimizer.step()
|
|
|
|
running_loss += loss.item()
|
|
total += labels.size(0)
|
|
_,predicted = torch.max(outputs.data,1)
|
|
#print(predicted)
|
|
#print("label",labels)
|
|
correct += (predicted == labels).sum().item()
|
|
train_acc = correct / total
|
|
|
|
print(f"[ Train | {epoch + 1:03d}/{num_epoch:03d} ] loss = {running_loss:.5f}, acc = {train_acc:.5f}")
|
|
|
|
##Validation
|
|
model.eval()
|
|
|
|
valid_loss = 0.0
|
|
total = 0
|
|
correct = 0
|
|
for i, data in enumerate(valid_loader):
|
|
inputs, labels = data
|
|
inputs = inputs.to(device)
|
|
labels = labels.to(device)
|
|
|
|
with torch.no_grad():
|
|
outputs = model(inputs)
|
|
|
|
loss = criterion(outputs, labels)
|
|
running_loss += loss.item()
|
|
total += labels.size(0)
|
|
_,predicted = torch.max(outputs.data,1)
|
|
correct += (predicted == labels).sum().item()
|
|
|
|
valid_acc = correct / total
|
|
print(f"[ Valid | {epoch + 1:03d}/{num_epoch:03d} ] loss = {running_loss:.5f}, acc = {valid_acc:.5f}")
|
|
if valid_acc > best_acc:
|
|
best_acc = valid_acc
|
|
torch.save(model.state_dict(), model_path)
|
|
print('saving model with acc {:.3f}'.format(valid_acc))
|
|
|
|
|
|
|
|
##Testing
|
|
model = models.resnet50()
|
|
model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3,
|
|
bias=False)
|
|
model.fc = nn.Linear(2048, 8)
|
|
model = model.to(device)
|
|
model.load_state_dict(torch.load(model_path))
|
|
model.eval()
|
|
|
|
with torch.no_grad():
|
|
correct = 0
|
|
total = 0
|
|
|
|
for i, data in enumerate(test_loader):
|
|
inputs, labels = data
|
|
|
|
inputs = inputs.to(device)
|
|
labels = labels.to(device)
|
|
|
|
outputs = model(inputs)
|
|
_,predicted = torch.max(outputs.data,1)
|
|
|
|
total += labels.size(0)
|
|
correct += (predicted == labels).sum().item()
|
|
# for k in range(batch_size):
|
|
# if predicted[k] != labels[k]:
|
|
# print(inputs[k])
|
|
|
|
#print(predicted)
|
|
#print("labels:",labels)
|
|
print('Test Accuracy:{} %'.format((correct / total) * 100))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|