2021-1-capstone-design2 / 2015103278

Go to a project
Toggle navigation Toggle navigation pinning
이동주
016f4d1e 1 parent e28e715f

1.train model with sampling data

Inline Side-by-side
Showing 1 changed file with 271 additions and 0 deletions
# run train.py --dataset cifar10 --model resnet18 --data_augmentation --cutout --length 16
# run train.py --dataset cifar100 --model resnet18 --data_augmentation --cutout --length 8
# run train.py --dataset svhn --model wideresnet --learning_rate 0.01 --epochs 160 --cutout --length 20
import pdb
import argparse
import numpy as np
from tqdm import tqdm
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.backends.cudnn as cudnn
from torch.optim.lr_scheduler import MultiStepLR
from torchvision.utils import make_grid, save_image
from torchvision import datasets, transforms
from torch.utils.data.dataloader import RandomSampler
from util.misc import CSVLogger
from util.cutout import Cutout
from model.resnet import ResNet18
from model.wide_resnet import WideResNet
model_options = ['resnet18', 'wideresnet']
dataset_options = ['cifar10', 'cifar100', 'svhn']
parser = argparse.ArgumentParser(description='CNN')
parser.add_argument('--dataset', '-d', default='cifar10',
choices=dataset_options)
parser.add_argument('--model', '-a', default='resnet18',
choices=model_options)
parser.add_argument('--batch_size', type=int, default=100,
help='input batch size for training (default: 128)')
parser.add_argument('--epochs', type=int, default=200,
help='number of epochs to train (default: 20)')
parser.add_argument('--learning_rate', type=float, default=0.1,
help='learning rate')
parser.add_argument('--data_augmentation', action='store_true', default=False,
help='augment data by flipping and cropping')
parser.add_argument('--cutout', action='store_true', default=False,
help='apply cutout')
parser.add_argument('--n_holes', type=int, default=0,
help='number of holes to cut out from image')
parser.add_argument('--length', type=int, default=0,
help='length of the holes')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='enables CUDA training')
parser.add_argument('--seed', type=int, default=0,
help='random seed (default: 1)')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
cudnn.benchmark = True # Should make training should go faster for large models
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
test_id = args.dataset + '_' + args.model
print(args)
# Image Preprocessing
if args.dataset == 'svhn':
normalize = transforms.Normalize(mean=[x / 255.0 for x in[109.9, 109.7, 113.8]],
std=[x / 255.0 for x in [50.1, 50.6, 50.8]])
else:
normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],
std=[x / 255.0 for x in [63.0, 62.1, 66.7]])
train_transform = transforms.Compose([])
if args.data_augmentation:
train_transform.transforms.append(transforms.RandomCrop(32, padding=4))
train_transform.transforms.append(transforms.RandomHorizontalFlip())
train_transform.transforms.append(transforms.ToTensor())
train_transform.transforms.append(normalize)
if args.cutout:
train_transform.transforms.append(Cutout(n_holes=args.n_holes, length=args.length))
test_transform = transforms.Compose([
transforms.ToTensor(),
normalize])
if args.dataset == 'cifar10':
num_classes = 10
train_dataset = datasets.CIFAR10(root='data/',
train=True,
transform=train_transform,
download=True)
test_dataset = datasets.CIFAR10(root='data/',
train=False,
transform=test_transform,
download=True)
elif args.dataset == 'cifar100':
num_classes = 100
train_dataset = datasets.CIFAR100(root='data/',
train=True,
transform=train_transform,
download=True)
test_dataset = datasets.CIFAR100(root='data/',
train=False,
transform=test_transform,
download=True)
elif args.dataset == 'svhn':
num_classes = 10
train_dataset = datasets.SVHN(root='data/',
split='train',
transform=train_transform,
download=True)
extra_dataset = datasets.SVHN(root='data/',
split='extra',
transform=train_transform,
download=True)
# Combine both training splits (https://arxiv.org/pdf/1605.07146.pdf)
data = np.concatenate([train_dataset.data, extra_dataset.data], axis=0)
labels = np.concatenate([train_dataset.labels, extra_dataset.labels], axis=0)
train_dataset.data = data
train_dataset.labels = labels
test_dataset = datasets.SVHN(root='data/',
split='test',
transform=test_transform,
download=True)
# Data Loader (Input Pipeline)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=False,
sampler=RandomSampler(train_dataset, True, 40000),
pin_memory=True,
num_workers=0)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=0)
if args.model == 'resnet18':
cnn = ResNet18(num_classes=num_classes)
elif args.model == 'wideresnet':
if args.dataset == 'svhn':
cnn = WideResNet(depth=16, num_classes=num_classes, widen_factor=8,
dropRate=0.4)
else:
cnn = WideResNet(depth=28, num_classes=num_classes, widen_factor=10,
dropRate=0.3)
cnn = cnn.cuda()
criterion = nn.CrossEntropyLoss().cuda()
cnn_optimizer = torch.optim.SGD(cnn.parameters(), lr=args.learning_rate,
momentum=0.9, nesterov=True, weight_decay=5e-4)
if args.dataset == 'svhn':
scheduler = MultiStepLR(cnn_optimizer, milestones=[80, 120], gamma=0.1)
else:
scheduler = MultiStepLR(cnn_optimizer, milestones=[60, 120, 160], gamma=0.2)
filename = 'logs/' + test_id + '.csv'
csv_logger = CSVLogger(args=args, fieldnames=['epoch', 'train_acc', 'test_acc', 'labels'], filename=filename)
def test(loader):
cnn.eval() # Change model to 'eval' mode (BN uses moving mean/var).
correct = 0.
total = 0.
count = 0
for images, labels in loader:
images = images.cuda()
labels = labels.cuda()
with torch.no_grad():
pred = cnn(images)
# print(pred, labels)
pred = torch.max(pred.data, 1)[1]
print(pred, '\n', labels, (pred==labels).sum().item())
# if (pred == labels).sum().item():
# print('match')
# count +=1
total += labels.size(0)
correct += (pred == labels).sum().item()
print(correct)
val_acc = correct / total
cnn.train()
return val_acc
for epoch in range(args.epochs):
xentropy_loss_avg = 0.
correct = 0.
total = 0.
# kl_sum = 0
# pred_sum = torch.Tensor([0] * 10).detach().cuda()
label_list = []
progress_bar = tqdm(train_loader)
for i, (images, labels) in enumerate(progress_bar):
progress_bar.set_description('Epoch ' + str(epoch))
images = images.cuda()
labels = labels.cuda()
cnn.zero_grad()
pred = cnn(images)
xentropy_loss = criterion(pred, labels)
xentropy_loss.backward()
cnn_optimizer.step()
xentropy_loss_avg += xentropy_loss.item()
# Calculate running average of accuracy
pred = torch.max(pred.data, 1)[1]
total += labels.size(0)
correct += (pred == labels.data).sum().item()
accuracy = correct / total
# print(pred)
# y_hat : 모델별 예측값 --> pred
# y_bar : 예측값들 평균값 -- > pred / total
# labes.data : ground_truth
# pred_sum = torch.add(pred_sum, pred)
# y_bar = pred_sum / (i+1)
# kl = torch.nn.functional.kl_div(pred, y_bar)
# kl_sum += kl
# for문 추가안하면 epoch별 iter마다 xentropy_loss_avg값의 1/iter이 xentropy값으로 출력
# for문 추가하면 epoch 별 iter 마다 xentropy_loss_avg 값은 동일하나 xentropy값 출력이 x_l_avg 값의 1/10으로 출력
# for문 상관 없이 pred, labels 값은 동일하게 확인됨.
# for a in range(list(pred_sum.size())[0]):
# for b in range(list(pred.size())[0]):
# if pred[b] == a:
# pred_sum[a] += 1
# print('\n',i, ' ', xentropy_loss_avg)
progress_bar.set_postfix(
# y_hat = '%.5f' % pred,
# y_bar = '%.5f' % y_bar,
# groun_truth = '%.5f' % labels.data,
# kl = '%.3f' % kl.item(),
# kl_sum = '%.3f' % (kl_sum.item()),
# kl_div = '%.3f' % (kl_sum.item() / (i + 1)), # kl_div 호출
xentropy='%.3f' % (xentropy_loss_avg / (i + 1)),
acc='%.3f' % accuracy)
# pred_sum = [x / 40000 for x in pred_sum]
test_acc = test(test_loader)
# print(pred, labels.data)
tqdm.write('test_acc: %.3f' % (test_acc))
scheduler.step(epoch) # Use this line for PyTorch <1.4
# scheduler.step() # Use this line for PyTorch >=1.4
row = {'epoch': str(epoch), 'train_acc': str(accuracy), 'test_acc': str(test_acc)}
csv_logger.writerow(row)
for i in range(len(label_list)):
csv_logger.writerow({'labels' : float(label_list[i])})
torch.save(cnn.state_dict(), 'checkpoints/' + test_id + '.pt')
csv_logger.close()