MNCJihun's picture
init
25322fb
import os
import sys
import matplotlib.pyplot as plt
from pandas.core.common import flatten
import torch
from torch import nn
from torch import optim
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms, models
import albumentations as A
from albumentations.pytorch import ToTensorV2
from tqdm import tqdm
import random
sys.path.append('/workspace')
import dataset
train_transforms = A.Compose(
[
A.SmallestMaxSize(max_size=350),
A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=360, p=0.5),
A.RandomCrop(height=256, width=256),
A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.5),
A.RandomBrightnessContrast(p=0.5),
A.MultiplicativeNoise(multiplier=[0.5,2], per_channel=True, p=0.2),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),
A.RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=0.5),
ToTensorV2(),
]
)
test_transforms = A.Compose(
[
A.SmallestMaxSize(max_size=350),
A.CenterCrop(height=256, width=256),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2(),
]
)
dataset_CV = dataset.MotorbikeDataset_CV(
root='/workspace/data/',
train_transforms=train_transforms,
val_transforms=test_transforms
)
train_dataset, val_dataset = dataset_CV.get_split()
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, drop_last=True)
val_loader = DataLoader(val_dataset,batch_size=64, shuffle=False)
device = torch.device("cuda:3") if torch.cuda.is_available() else torch.device("cpu")
model = models.resnet50(pretrained=True)
model.fc = nn.Sequential(
nn.Dropout(0.5),
nn.Linear(model.fc.in_features, 2)
)
for n, p in model.named_parameters():
if 'fc' in n:
p.requires_grad = True
else:
p.requires_grad = False
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
best_acc = 0.0
for epoch in range(10):
model.train()
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
inputs, labels = data[0].to(device), data[1].to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
scheduler.step()
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
# print("TRAIN acc = {}".format(acc))
running_loss = 0.0
with torch.no_grad():
model.eval()
running_loss = 0.0
correct =0
for i, data in enumerate(val_loader, 0):
inputs, labels = data[0].to(device), data[1].to(device)
outputs = model(inputs)
_, preds = outputs.max(1)
loss = criterion(outputs, labels)
running_loss += loss.item()
labels_one_hot = F.one_hot(labels, 2)
outputs_one_hot = F.one_hot(preds, 2)
correct = correct + (labels_one_hot + outputs_one_hot == 2).sum().to(torch.float)
acc = 100 * correct / len(val_dataset)
print(f'VAL: [{epoch + 1}, {i + 1:5d}] loss: {running_loss / len(val_loader):.3f}')
print("VAL acc = {:.2f}".format(acc))
if best_acc < acc:
torch.save(model.state_dict(), './result/best_model.pth')