Created
March 27, 2017 10:58
-
-
Save FuriouslyCurious/698e95a4c79fdaa37c32d3b2076a95ac to your computer and use it in GitHub Desktop.
From Brandon Amos's Densenet Implementation - Data
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python3 | |
import argparse | |
import torch | |
import torch.nn as nn | |
import torch.optim as optim | |
import torch.nn.functional as F | |
from torch.autograd import Variable | |
import torchvision.datasets as dset | |
import torchvision.transforms as transforms | |
from torchvision.utils import save_image | |
from torch.utils.data import DataLoader | |
import os | |
import sys | |
import math | |
import shutil | |
import setproctitle | |
import densenet | |
#import make_graph | |
def main(): | |
parser = argparse.ArgumentParser() | |
parser.add_argument('--batchSz', type=int, default=32) | |
parser.add_argument('--nEpochs', type=int, default=11) | |
parser.add_argument('--no-cuda', action='store_false') | |
parser.add_argument('--save') | |
parser.add_argument('--seed', type=int, default=1) | |
parser.add_argument('--opt', type=str, default='adam', | |
choices=('sgd', 'adam', 'rmsprop')) | |
args = parser.parse_args() | |
args.cuda = True #not args.no_cuda and torch.cuda.is_available() | |
args.save = args.save or 'work/densenet.base' | |
setproctitle.setproctitle(args.save) | |
torch.manual_seed(args.seed) | |
if args.cuda: | |
torch.cuda.manual_seed(args.seed) | |
if os.path.exists(args.save): | |
shutil.rmtree(args.save) | |
os.makedirs(args.save, exist_ok=True) | |
normMean = [0.6492440333064673, 0.6492440333064673, 0.6492440333064673] | |
normStd = [0.235452743560426, 0.235452743560426, 0.235452743560426] | |
normTransform = transforms.Normalize(normMean, normStd) | |
trainTransform = transforms.Compose([ | |
#transforms.RandomCrop(32, padding=4), | |
transforms.RandomHorizontalFlip(), | |
transforms.ToTensor(), | |
normTransform | |
]) | |
testTransform = transforms.Compose([ | |
transforms.ToTensor(), | |
normTransform | |
]) | |
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} | |
trainLoader = DataLoader( | |
dset.ImageFolder(root='/home/FC/data/P/train', # train=True, download=True, | |
transform=trainTransform), | |
batch_size=args.batchSz, shuffle=True, **kwargs) | |
testLoader = DataLoader( | |
dset.ImageFolder(root='/home/FC/data/P/val', # train=False, download=True, | |
transform=testTransform), | |
batch_size=args.batchSz, shuffle=False, **kwargs) | |
net = densenet.DenseNet(growthRate=12, depth=100, reduction=0.5, | |
bottleneck=True, nClasses=3) | |
print(' + Number of params: {}'.format( | |
sum([p.data.nelement() for p in net.parameters()]))) | |
if args.cuda: | |
net = net.cuda() | |
if args.opt == 'sgd': | |
optimizer = optim.SGD(net.parameters(), lr=1e-1, | |
momentum=0.9, weight_decay=1e-4) | |
elif args.opt == 'adam': | |
optimizer = optim.Adam(net.parameters(), weight_decay=1e-4) | |
elif args.opt == 'rmsprop': | |
optimizer = optim.RMSprop(net.parameters(), weight_decay=1e-4) | |
trainF = open(os.path.join(args.save, 'train.csv'), 'w') | |
testF = open(os.path.join(args.save, 'test.csv'), 'w') | |
for epoch in range(1, args.nEpochs + 1): | |
adjust_opt(args.opt, optimizer, epoch) | |
train(args, epoch, net, trainLoader, optimizer, trainF) | |
test(args, epoch, net, testLoader, optimizer, testF) | |
torch.save(net, os.path.join(args.save, 'latest.pth')) | |
os.system('./plot.py {} &'.format(args.save)) | |
trainF.close() | |
testF.close() | |
def train(args, epoch, net, trainLoader, optimizer, trainF): | |
net.train() | |
nProcessed = 0 | |
nTrain = len(trainLoader.dataset) | |
for batch_idx, (data, target) in enumerate(trainLoader): | |
if args.cuda: | |
data, target = data.cuda(), target.cuda() | |
data, target = Variable(data), Variable(target) | |
optimizer.zero_grad() | |
output = net(data) | |
loss = F.nll_loss(output, target) | |
# make_graph.save('/tmp/t.dot', loss.creator); assert(False) | |
loss.backward() | |
optimizer.step() | |
nProcessed += len(data) | |
pred = output.data.max(1)[1] # get the index of the max log-probability | |
incorrect = pred.ne(target.data).cpu().sum() | |
err = 100.*incorrect/len(data) | |
partialEpoch = epoch + batch_idx / len(trainLoader) - 1 | |
print('Train Epoch: {:.2f} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tError: {:.6f}'.format( | |
partialEpoch, nProcessed, nTrain, 100. * batch_idx / len(trainLoader), | |
loss.data[0], err)) | |
trainF.write('{},{},{}\n'.format(partialEpoch, loss.data[0], err)) | |
trainF.flush() | |
def test(args, epoch, net, testLoader, optimizer, testF): | |
net.eval() | |
test_loss = 0 | |
incorrect = 0 | |
for data, target in testLoader: | |
if args.cuda: | |
data, target = data.cuda(), target.cuda() | |
data, target = Variable(data, volatile=True), Variable(target) | |
output = net(data) | |
test_loss += F.nll_loss(output, target).data[0] | |
pred = output.data.max(1)[1] # get the index of the max log-probability | |
incorrect += pred.ne(target.data).cpu().sum() | |
test_loss = test_loss | |
test_loss /= len(testLoader) # loss function already averages over batch size | |
nTotal = len(testLoader.dataset) | |
err = 100.*incorrect/nTotal | |
print('\nTest set: Average loss: {:.4f}, Error: {}/{} ({:.0f}%)\n'.format( | |
test_loss, incorrect, nTotal, err)) | |
testF.write('{},{},{}\n'.format(epoch, test_loss, err)) | |
testF.flush() | |
def adjust_opt(optAlg, optimizer, epoch): | |
if optAlg == 'sgd': | |
if epoch < 150: lr = 1e-1 | |
elif epoch == 150: lr = 1e-2 | |
elif epoch == 225: lr = 1e-3 | |
else: return | |
for param_group in optimizer.param_groups: | |
param_group['lr'] = lr | |
if __name__=='__main__': | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment