Skip to content

GitLab

Unverified Commit b872eb8c authored by Hang Zhang's avatar Hang Zhang Committed by GitHub
Browse files

ResNeSt plus (#256)

parent 5a1e3fbc
Hide whitespace changes
Inline Side-by-side
Showing with 1254 additions and 1032 deletions
experiments/recognition/dataset/cifar10.py deleted 100644 → 0
View file @ 5a1e3fbc
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import torch
import torchvision
import torchvision.transforms as transforms
class Dataloader():
def __init__(self, args):
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform_train)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform_test)
kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
trainloader = torch.utils.data.DataLoader(trainset, batch_size=
args.batch_size, shuffle=True, **kwargs)
testloader = torch.utils.data.DataLoader(testset, batch_size=
args.batch_size, shuffle=False, **kwargs)
self.trainloader = trainloader
self.testloader = testloader
def getloader(self):
return self.trainloader, self.testloader
experiments/recognition/dataset/minc.py deleted 100644 → 0
View file @ 5a1e3fbc
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import torch
import torch.utils.data as data
import torchvision
from torchvision import transforms
from PIL import Image
import os
import os.path
_imagenet_pca = {
'eigval': torch.Tensor([0.2175, 0.0188, 0.0045]),
'eigvec': torch.Tensor([
[-0.5675, 0.7192, 0.4009],
[-0.5808, -0.0045, -0.8140],
[-0.5836, -0.6948, 0.4203],
])
}
def find_classes(dir):
classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
classes.sort()
class_to_idx = {classes[i]: i for i in range(len(classes))}
return classes, class_to_idx
def make_dataset(filename, datadir, class_to_idx):
images = []
labels = []
with open(os.path.join(filename), "r") as lines:
for line in lines:
_image = os.path.join(datadir, line.rstrip('\n'))
_dirname = os.path.split(os.path.dirname(_image))[1]
assert os.path.isfile(_image)
label = class_to_idx[_dirname]
images.append(_image)
labels.append(label)
return images, labels
class MINCDataset(data.Dataset):
def __init__(self, root, train=True, transform=None):
self.transform = transform
classes, class_to_idx = find_classes(root + '/images')
if train:
filename = os.path.join(root, 'labels/train1.txt')
else:
filename = os.path.join(root, 'labels/test1.txt')
self.images, self.labels = make_dataset(filename, root,
class_to_idx)
assert (len(self.images) == len(self.labels))
def __getitem__(self, index):
_img = Image.open(self.images[index]).convert('RGB')
_label = self.labels[index]
if self.transform is not None:
_img = self.transform(_img)
return _img, _label
def __len__(self):
return len(self.images)
class Dataloader():
def __init__(self, args):
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform_train = transforms.Compose([
transforms.Resize(256),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(0.4,0.4,0.4),
transforms.ToTensor(),
Lighting(0.1, _imagenet_pca['eigval'], _imagenet_pca['eigvec']),
normalize,
])
transform_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])
trainset = MINCDataset(root=os.path.expanduser('~/.encoding/data/minc-2500/'),
train=True, transform=transform_train)
testset = MINCDataset(root=os.path.expanduser('~/.encoding/data/minc-2500/'),
train=False, transform=transform_test)
kwargs = {'num_workers': 8, 'pin_memory': True} if args.cuda else {}
trainloader = torch.utils.data.DataLoader(trainset, batch_size=
args.batch_size, shuffle=True, **kwargs)
testloader = torch.utils.data.DataLoader(testset, batch_size=
args.test_batch_size, shuffle=False, **kwargs)
self.trainloader = trainloader
self.testloader = testloader
def getloader(self):
return self.trainloader, self.testloader
class Lighting(object):
"""Lighting noise(AlexNet - style PCA - based noise)"""
def __init__(self, alphastd, eigval, eigvec):
self.alphastd = alphastd
self.eigval = eigval
self.eigvec = eigvec
def __call__(self, img):
if self.alphastd == 0:
return img
alpha = img.new().resize_(3).normal_(0, self.alphastd)
rgb = self.eigvec.type_as(img).clone()\
.mul(alpha.view(1, 3).expand(3, 3))\
.mul(self.eigval.view(1, 3).expand(3, 3))\
.sum(1).squeeze()
return img.add(rgb.view(3, 1, 1).expand_as(img))
if __name__ == "__main__":
trainset = MINCDataset(root=os.path.expanduser('~/.encoding/data/minc-2500/'), train=True)
testset = MINCDataset(root=os.path.expanduser('~/.encoding/data/minc-2500/'), train=False)
print(len(trainset))
print(len(testset))
experiments/recognition/main.py
View file @ b872eb8c
......@@ -10,13 +10,94 @@
from __future__ import print_function
import os
import argparse
from tqdm import tqdm
import torch
import torch.nn as nn
import encoding
from option import Options
from encoding.nn import LabelSmoothing, NLLMultiLabelSmooth
from encoding.utils import (accuracy, AverageMeter, MixUpWrapper, LR_Scheduler)
class Options():
def __init__(self):
# data settings
parser = argparse.ArgumentParser(description='Deep Encoding')
parser.add_argument('--dataset', type=str, default='cifar10',
help='training dataset (default: cifar10)')
parser.add_argument('--base-size', type=int, default=None,
help='base image size')
parser.add_argument('--crop-size', type=int, default=224,
help='crop image size')
parser.add_argument('--label-smoothing', type=float, default=0.0,
help='label-smoothing (default eta: 0.0)')
parser.add_argument('--mixup', type=float, default=0.0,
help='mixup (default eta: 0.0)')
parser.add_argument('--rand-aug', action='store_true',
default=False, help='rectify convolution')
# model params
parser.add_argument('--model', type=str, default='densenet',
help='network model type (default: densenet)')
parser.add_argument('--pretrained', action='store_true',
default=False, help='load pretrianed mode')
parser.add_argument('--rectify', action='store_true',
default=False, help='rectify convolution')
parser.add_argument('--rectify-avg', action='store_true',
default=False, help='rectify convolution')
parser.add_argument('--last-gamma', action='store_true', default=False,
help='whether to init gamma of the last BN layer in \
each bottleneck to 0 (default: False)')
parser.add_argument('--dropblock-prob', type=float, default=0,
help='DropBlock prob. default is 0.')
parser.add_argument('--final-drop', type=float, default=0,
help='final dropout prob. default is 0.')
# training hyper params
parser.add_argument('--batch-size', type=int, default=128, metavar='N',
help='batch size for training (default: 128)')
parser.add_argument('--test-batch-size', type=int, default=256, metavar='N',
help='batch size for testing (default: 256)')
parser.add_argument('--epochs', type=int, default=120, metavar='N',
help='number of epochs to train (default: 600)')
parser.add_argument('--start_epoch', type=int, default=0,
metavar='N', help='the epoch number to start (default: 1)')
parser.add_argument('--workers', type=int, default=32,
metavar='N', help='dataloader threads')
# optimizer
parser.add_argument('--lr', type=float, default=0.1, metavar='LR',
help='learning rate (default: 0.1)')
parser.add_argument('--lr-scheduler', type=str, default='cos',
help='learning rate scheduler (default: cos)')
parser.add_argument('--warmup-epochs', type=int, default=0,
help='number of warmup epochs (default: 0)')
parser.add_argument('--lr-step', type=int, default=40, metavar='LR',
help='learning rate step (default: 40)')
parser.add_argument('--momentum', type=float, default=0.9,
metavar='M', help='SGD momentum (default: 0.9)')
parser.add_argument('--weight-decay', type=float, default=1e-4,
metavar ='M', help='SGD weight decay (default: 1e-4)')
parser.add_argument('--no-bn-wd', action='store_true',
default=False, help='no bias decay')
# cuda, seed and logging
parser.add_argument('--no-cuda', action='store_true',
default=False, help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
# checking point
parser.add_argument('--resume', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--checkname', type=str, default='default',
help='set the checkpoint name')
# evaluation option
parser.add_argument('--eval', action='store_true', default= False,
help='evaluating')
parser.add_argument('--export', type=str, default=None,
help='put the path to resuming file if needed')
self.parser = parser
def parse(self):
args = self.parser.parse_args()
return args
# global variable
best_pred = 0.0
......@@ -33,38 +114,87 @@ def main():
if args.cuda:
torch.cuda.manual_seed(args.seed)
# init dataloader
transform_train, transform_val = encoding.transforms.get_transform(args.dataset)
transform_train, transform_val = encoding.transforms.get_transform(
args.dataset, args.base_size, args.crop_size, args.rand_aug)
trainset = encoding.datasets.get_dataset(args.dataset, root=os.path.expanduser('~/.encoding/data'),
transform=transform_train, train=True, download=True)
valset = encoding.datasets.get_dataset(args.dataset, root=os.path.expanduser('~/.encoding/data'),
transform=transform_val, train=False, download=True)
train_loader = torch.utils.data.DataLoader(
trainset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
num_workers=args.workers, drop_last=True, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
valset, batch_size=args.test_batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
# init the model
model = encoding.models.get_model(args.model, pretrained=args.pretrained)
model_kwargs = {}
if args.pretrained:
model_kwargs['pretrained'] = True
if args.final_drop > 0.0:
model_kwargs['final_drop'] = args.final_drop
if args.dropblock_prob > 0.0:
model_kwargs['dropblock_prob'] = args.dropblock_prob
if args.last_gamma:
model_kwargs['last_gamma'] = True
if args.rectify:
model_kwargs['rectified_conv'] = True
model_kwargs['rectify_avg'] = args.rectify_avg
model = encoding.models.get_model(args.model, **model_kwargs)
if args.dropblock_prob > 0.0:
from functools import partial
from encoding.nn import reset_dropblock
nr_iters = (args.epochs - 2 * args.warmup_epochs) * len(train_loader)
apply_drop_prob = partial(reset_dropblock, args.warmup_epochs*len(train_loader),
nr_iters, 0.0, args.dropblock_prob)
model.apply(apply_drop_prob)
print(model)
# criterion and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.mixup > 0:
train_loader = MixUpWrapper(args.mixup, 1000, train_loader,
list(range(torch.cuda.device_count())))
criterion = NLLMultiLabelSmooth(args.label_smoothing)
elif args.label_smoothing > 0.0:
criterion = LabelSmoothing(args.label_smoothing)
else:
criterion = nn.CrossEntropyLoss()
if args.no_bn_wd:
parameters = model.named_parameters()
param_dict = {}
for k, v in parameters:
param_dict[k] = v
bn_params = [v for n, v in param_dict.items() if ('bn' in n or 'bias' in n)]
rest_params = [v for n, v in param_dict.items() if not ('bn' in n or 'bias' in n)]
print(" Weight decay NOT applied to BN parameters ")
print(f'len(parameters): {len(list(model.parameters()))} = {len(bn_params)} + {len(rest_params)}')
optimizer = torch.optim.SGD([{'params': bn_params, 'weight_decay': 0 },
{'params': rest_params, 'weight_decay': args.weight_decay}],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.cuda:
model.cuda()
criterion.cuda()
# Please use CUDA_VISIBLE_DEVICES to control the number of gpus
model = nn.DataParallel(model)
# check point
if args.resume is not None:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] +1
args.start_epoch = checkpoint['epoch'] + 1 if args.start_epoch == 1 else args.start_epoch
best_pred = checkpoint['best_pred']
acclist_train = checkpoint['acclist_train']
acclist_val = checkpoint['acclist_val']
......@@ -75,8 +205,12 @@ def main():
else:
raise RuntimeError ("=> no resume checkpoint found at '{}'".\
format(args.resume))
scheduler = encoding.utils.LR_Scheduler(args.lr_scheduler, args.lr, args.epochs,
len(train_loader), args.lr_step)
scheduler = LR_Scheduler(args.lr_scheduler,
base_lr=args.lr,
num_epochs=args.epochs,
iters_per_epoch=len(train_loader),
warmup_epochs=args.warmup_epochs,
lr_step=args.lr_step)
def train(epoch):
model.train()
losses = AverageMeter()
......@@ -85,6 +219,7 @@ def main():
tbar = tqdm(train_loader, desc='\r')
for batch_idx, (data, target) in enumerate(tbar):
scheduler(optimizer, batch_idx, epoch, best_pred)
#criterion.update(batch_idx, epoch)
if args.cuda:
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
......@@ -127,6 +262,7 @@ def main():
best_pred = top1.avg
is_best = True
encoding.utils.save_checkpoint({
'args': args,
'epoch': epoch,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
......@@ -135,47 +271,19 @@ def main():
'acclist_val':acclist_val,
}, args=args, is_best=is_best)
if args.export:
torch.save(model.module.state_dict(), args.export + '.pth')
return
if args.eval:
validate(args.start_epoch)
return
for epoch in range(args.start_epoch, args.epochs + 1):
for epoch in range(args.start_epoch, args.epochs):
train(epoch)
validate(epoch)
def accuracy(output, target, topk=(1,)):
"""Computes the accuracy over the k top predictions for the specified values of k"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(100.0 / batch_size))
return res
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
validate(epoch)
if __name__ == "__main__":
main()
experiments/recognition/model/deepten.py deleted 100644 → 0
View file @ 5a1e3fbc
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import torch
from torch.autograd import Variable
import torch.nn as nn
from torch.autograd import Variable
import encoding
import encoding.models.resnet as resnet
class Net(nn.Module):
def __init__(self, args):
nclass=args.nclass
super(Net, self).__init__()
self.backbone = args.backbone
# copying modules from pretrained models
if self.backbone == 'resnet50':
self.pretrained = resnet.resnet50(pretrained=True, dilated=False)
elif self.backbone == 'resnet101':
self.pretrained = resnet.resnet101(pretrained=True, dilated=False)
elif self.backbone == 'resnet152':
self.pretrained = resnet.resnet152(pretrained=True, dilated=False)
else:
raise RuntimeError('unknown backbone: {}'.format(self.backbone))
n_codes = 32
self.head = nn.Sequential(
nn.Conv2d(2048, 128, 1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
encoding.nn.Encoding(D=128,K=n_codes),
encoding.nn.View(-1, 128*n_codes),
encoding.nn.Normalize(),
nn.Linear(128*n_codes, nclass),
)
def forward(self, x):
if isinstance(x, Variable):
_, _, h, w = x.size()
elif isinstance(x, tuple) or isinstance(x, list):
var_input = x
while not isinstance(var_input, Variable):
var_input = var_input[0]
_, _, h, w = var_input.size()
else:
raise RuntimeError('unknown input type: ', type(x))
x = self.pretrained.conv1(x)
x = self.pretrained.bn1(x)
x = self.pretrained.relu(x)
x = self.pretrained.maxpool(x)
x = self.pretrained.layer1(x)
x = self.pretrained.layer2(x)
x = self.pretrained.layer3(x)
x = self.pretrained.layer4(x)
return self.head(x)
experiments/recognition/model/encnet.py deleted 100644 → 0
View file @ 5a1e3fbc
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import torch
from torch.autograd import Variable
import torch.nn as nn
from .mynn import EncBasicBlock
import encoding
class Net(nn.Module):
def __init__(self, args):
super(Net, self).__init__()
num_blocks=[2,2,2]
block=EncBasicBlock
if block == EncBasicBlock:
self.expansion = 1
else:
self.expansion = 4
self.inplanes = args.widen * 16
strides = [1, 2, 2]
model = []
# Conv_1
model += [nn.Conv2d(3, self.inplanes, kernel_size=3, padding=1),
nn.BatchNorm2d(self.inplanes),
nn.ReLU(inplace=True)]
# Residual units
model += [self._residual_unit(block, self.inplanes, num_blocks[0],
strides[0], args.ncodes)]
for i in range(2):
model += [self._residual_unit(block,
int(2*self.inplanes/self.expansion),
num_blocks[i+1], strides[i+1], args.ncodes)]
# Last conv layer
model += [nn.BatchNorm2d(self.inplanes),
nn.ReLU(inplace=True),
nn.AvgPool2d(8),
encoding.nn.View(-1, self.inplanes),
nn.Linear(self.inplanes, args.nclass)]
self.model = nn.Sequential(*model)
def _residual_unit(self, block, planes, n_blocks, stride, ncodes):
strides = [stride] + [1]*(n_blocks-1)
layers = []
for i in range(n_blocks):
layers += [block(self.inplanes, planes, strides[i], ncodes)]
self.inplanes = self.expansion*planes
return nn.Sequential(*layers)
def forward(self, input):
return self.model(input)
experiments/recognition/model/encnetdrop.py deleted 100644 → 0
View file @ 5a1e3fbc
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import torch
from torch.autograd import Variable
import torch.nn as nn
from .mynn import EncBasicBlock, EncDropLayer
import encoding
class Net(nn.Module):
def __init__(self, args):
super(Net, self).__init__()
num_blocks=[2,2,2]
block=EncBasicBlock
if block == EncBasicBlock:
self.expansion = 1
else:
self.expansion = 4
self.inplanes = args.widen * 16
strides = [1, 2, 2]
model = []
# Conv_1
model += [nn.Conv2d(3, self.inplanes, kernel_size=3, padding=1),
nn.BatchNorm2d(self.inplanes),
nn.ReLU(inplace=True)]
# Residual units
model += [self._residual_unit(block, self.inplanes, num_blocks[0],
strides[0], args.ncodes)]
for i in range(2):
model += [self._residual_unit(block,
int(2*self.inplanes/self.expansion),
num_blocks[i+1], strides[i+1], args.ncodes)]
# Last conv layer
model += [nn.BatchNorm2d(self.inplanes),
nn.ReLU(inplace=True),
nn.AvgPool2d(8),
encoding.nn.View(-1, self.inplanes),
nn.Linear(self.inplanes, args.nclass)]
self.model = nn.Sequential(*model)
def _residual_unit(self, block, planes, n_blocks, stride, ncodes):
strides = [stride] + [1]*(n_blocks-1)
layers = []
for i in range(n_blocks):
layers += [block(self.inplanes, planes, strides[i], ncodes, ELayer=EncDropLayer)]
self.inplanes = self.expansion*planes
return nn.Sequential(*layers)
def forward(self, input):
return self.model(input)
def test():
net = Net().cuda()
print(net)
x = Variable(torch.randn(1,3,32,32)).cuda()
y = net(x)
print(y)
params = net.parameters()
sum = 0
for param in params:
sum += param.nelement()
print('Total params:', sum)
if __name__ == "__main__":
test()
experiments/recognition/model/mynn.py deleted 100644 → 0
View file @ 5a1e3fbc
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import encoding
def conv3x3(in_planes, out_planes, stride=1):
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
class Basicblock(nn.Module):
""" Pre-activation residual block
Identity Mapping in Deep Residual Networks
ref https://arxiv.org/abs/1603.05027
"""
def __init__(self, inplanes, planes, stride=1,
norm_layer=nn.BatchNorm2d):
super(Basicblock, self).__init__()
if inplanes != planes or stride !=1 :
self.downsample = True
self.residual_layer = nn.Conv2d(inplanes, planes,
kernel_size=1, stride=stride)
else:
self.downsample = False
conv_block=[]
conv_block+=[norm_layer(inplanes),
nn.ReLU(inplace=True),
conv3x3(inplanes, planes,stride=stride),
norm_layer(planes),
nn.ReLU(inplace=True),
conv3x3(planes, planes)]
self.conv_block = nn.Sequential(*conv_block)
def forward(self, input):
if self.downsample:
residual = self.residual_layer(input)
else:
residual = input
return residual + self.conv_block(input)
class Bottleneck(nn.Module):
""" Pre-activation residual block
Identity Mapping in Deep Residual Networks
ref https://arxiv.org/abs/1603.05027
"""
def __init__(self, inplanes, planes, stride=1,norm_layer=nn.BatchNorm2d):
super(Bottleneck, self).__init__()
self.expansion = 4
if inplanes != planes*self.expansion or stride !=1 :
self.downsample = True
self.residual_layer = nn.Conv2d(inplanes,
planes * self.expansion, kernel_size=1, stride=stride)
else:
self.downsample = False
conv_block = []
conv_block += [norm_layer(inplanes),
nn.ReLU(inplace=True),
nn.Conv2d(inplanes, planes, kernel_size=1,
stride=1, bias=False)]
conv_block += [norm_layer(planes),
nn.ReLU(inplace=True),
nn.Conv2d(planes, planes, kernel_size=3,
stride=stride, padding=1, bias=False)]
conv_block += [norm_layer(planes),
nn.ReLU(inplace=True),
nn.Conv2d(planes, planes * self.expansion,
kernel_size=1, stride=1, bias=False)]
self.conv_block = nn.Sequential(*conv_block)
def forward(self, x):
if self.downsample:
residual = self.residual_layer(x)
else:
residual = x
return residual + self.conv_block(x)
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
class EncLayerV2(nn.Module):
def __init__(self, channel, K=16, reduction=4):
super(EncLayerV2, self).__init__()
out_channel = int(channel / reduction)
self.fc = nn.Sequential(
nn.Conv2d(channel, out_channel, 1),
nn.BatchNorm2d(out_channel),
nn.ReLU(inplace=True),
encoding.nn.EncodingV2(D=out_channel,K=K),
encoding.nn.View(-1, out_channel*K),
encoding.nn.Normalize(),
nn.Linear(out_channel*K, channel),
nn.Sigmoid()
)
def forward(self, x):
b, c, _, _ = x.size()
y = self.fc(x).view(b, c, 1, 1)
return x * y
class EncLayerV3(nn.Module):
def __init__(self, channel, K=16, reduction=4):
super(EncLayerV3, self).__init__()
out_channel = int(channel / reduction)
self.fc = nn.Sequential(
nn.Conv2d(channel, out_channel, 1),
nn.BatchNorm2d(out_channel),
nn.ReLU(inplace=True),
encoding.nn.EncodingV3(D=out_channel,K=K),
encoding.nn.View(-1, out_channel*K),
encoding.nn.Normalize(),
nn.Linear(out_channel*K, channel),
nn.Sigmoid()
)
def forward(self, x):
b, c, _, _ = x.size()
y = self.fc(x).view(b, c, 1, 1)
return x * y
class EncLayer(nn.Module):
def __init__(self, channel, K=16, reduction=4):
super(EncLayer, self).__init__()
out_channel = int(channel / reduction)
self.fc = nn.Sequential(
nn.Conv2d(channel, out_channel, 1),
nn.BatchNorm2d(out_channel),
nn.ReLU(inplace=True),
encoding.nn.Encoding(D=out_channel,K=K),
encoding.nn.View(-1, out_channel*K),
encoding.nn.Normalize(),
nn.Linear(out_channel*K, channel),
nn.Sigmoid()
)
def forward(self, x):
b, c, _, _ = x.size()
y = self.fc(x).view(b, c, 1, 1)
return x * y
class EncDropLayer(nn.Module):
def __init__(self, channel, K=16, reduction=4):
super(EncDropLayer, self).__init__()
out_channel = int(channel / reduction)
self.fc = nn.Sequential(
nn.Conv2d(channel, out_channel, 1),
nn.BatchNorm2d(out_channel),
nn.ReLU(inplace=True),
encoding.nn.EncodingDrop(D=out_channel,K=K),
encoding.nn.View(-1, out_channel*K),
encoding.nn.Normalize(),
nn.Linear(out_channel*K, channel),
nn.Sigmoid()
)
def forward(self, x):
b, c, _, _ = x.size()
y = self.fc(x).view(b, c, 1, 1)
return x * y
class EncBasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, K=16, ELayer=EncLayer):
super(EncBasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes, 1)
self.bn2 = nn.BatchNorm2d(planes)
self.se = ELayer(planes, K, self.expansion*4)
self.stride = stride
if inplanes != planes or stride !=1 :
self.downsample = True
self.residual_layer = nn.Conv2d(inplanes, planes,
kernel_size=1, stride=stride)
else:
self.downsample = False
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.se(out)
if self.downsample:
residual = self.residual_layer(x)
out += residual
out = self.relu(out)
return out
class EncBottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None, K=16, ELayer=EncLayer):
super(EncBottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
stride=stride, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = nn.Conv2d(planes, planes * self.expansion,
kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
self.relu = nn.ReLU(inplace=True)
self.se = ELayer(planes * self.expansion, K, self.expansion*4)
self.stride = stride
if inplanes != planes * self.expansion or stride !=1 :
self.downsample = True
self.residual_layer = nn.Conv2d(inplanes,
planes* self.expansion, kernel_size=1, stride=stride)
else:
self.downsample = False
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
out = self.se(out)
if self.downsample:
residual = self.residual_layer(x)
out += residual
out = self.relu(out)
return out
experiments/recognition/model/resnet.py deleted 100644 → 0
View file @ 5a1e3fbc
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import torch
import torch.nn as nn
from torch.autograd import Variable
import model.mynn as nn2
import encoding
class Net(nn.Module):
def __init__(self, args):
super(Net, self).__init__()
num_blocks=[2,2,2]
block=nn2.Basicblock
if block == nn2.Basicblock:
self.expansion = 1
else:
self.expansion = 4
self.inplanes = args.widen * 16
strides = [1, 2, 2]
model = []
# Conv_1
model += [nn.Conv2d(3, self.inplanes, kernel_size=3, padding=1),
nn.BatchNorm2d(self.inplanes),
nn.ReLU(inplace=True)]
# Residual units
model += [self._residual_unit(block, self.inplanes, num_blocks[0],
strides[0])]
for i in range(2):
model += [self._residual_unit(block,
int(2*self.inplanes/self.expansion), num_blocks[i+1],
strides[i+1])]
# Last conv layer
model += [nn.BatchNorm2d(self.inplanes),
nn.ReLU(inplace=True),
nn.AvgPool2d(8),
encoding.nn.View(-1, self.inplanes),
nn.Linear(self.inplanes, args.nclass)]
self.model = nn.Sequential(*model)
def _residual_unit(self, block, planes, n_blocks, stride):
strides = [stride] + [1]*(n_blocks-1)
layers = []
for i in range(n_blocks):
layers += [block(self.inplanes, planes, strides[i])]
self.inplanes = self.expansion*planes
return nn.Sequential(*layers)
def forward(self, input):
return self.model(input)
experiments/recognition/option.py deleted 100644 → 0
View file @ 5a1e3fbc
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import argparse
import os
class Options():
def __init__(self):
# Training settings
parser = argparse.ArgumentParser(description='Deep Encoding')
parser.add_argument('--dataset', type=str, default='cifar10',
help='training dataset (default: cifar10)')
# model params
parser.add_argument('--model', type=str, default='densenet',
help='network model type (default: densenet)')
parser.add_argument('--pretrained', action='store_true',
default=False, help='load pretrianed mode')
parser.add_argument('--nclass', type=int, default=10, metavar='N',
help='number of classes (default: 10)')
parser.add_argument('--widen', type=int, default=4, metavar='N',
help='widen factor of the network (default: 4)')
parser.add_argument('--ncodes', type=int, default=32, metavar='N',
help='number of codewords in Encoding Layer (default: 32)')
parser.add_argument('--backbone', type=str, default='resnet50',
help='backbone name (default: resnet50)')
# training hyper params
parser.add_argument('--batch-size', type=int, default=128,
metavar='N', help='batch size for training (default: 128)')
parser.add_argument('--test-batch-size', type=int, default=256,
metavar='N', help='batch size for testing (default: 256)')
parser.add_argument('--epochs', type=int, default=600, metavar='N',
help='number of epochs to train (default: 600)')
parser.add_argument('--start_epoch', type=int, default=1,
metavar='N', help='the epoch number to start (default: 1)')
parser.add_argument('--workers', type=int, default=16,
metavar='N', help='dataloader threads')
# lr setting
parser.add_argument('--lr', type=float, default=0.1, metavar='LR',
help='learning rate (default: 0.1)')
parser.add_argument('--lr-scheduler', type=str, default='cos',
help='learning rate scheduler (default: cos)')
parser.add_argument('--lr-step', type=int, default=40, metavar='LR',
help='learning rate step (default: 40)')
# optimizer
parser.add_argument('--momentum', type=float, default=0.9,
metavar='M', help='SGD momentum (default: 0.9)')
parser.add_argument('--weight-decay', type=float, default=1e-4,
metavar ='M', help='SGD weight decay (default: 1e-4)')
# cuda, seed and logging
parser.add_argument('--no-cuda', action='store_true',
default=False, help='disables CUDA training')
parser.add_argument('--plot', action='store_true', default=False,
help='matplotlib')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
# checking point
parser.add_argument('--resume', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--checkname', type=str, default='default',
help='set the checkpoint name')
# evaluation option
parser.add_argument('--eval', action='store_true', default= False,
help='evaluating')
self.parser = parser
def parse(self):
args = self.parser.parse_args()
if args.dataset == 'minc':
args.nclass = 23
return args
experiments/recognition/train_dist.py 0 → 100644
View file @ b872eb8c
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## Email: zhanghang0704@gmail.com
## Copyright (c) 2020
##
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import os
import time
import argparse
import numpy as np
from tqdm import tqdm
from mpi4py import MPI
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel
import encoding
from encoding.nn import LabelSmoothing, NLLMultiLabelSmooth
from encoding.utils import (accuracy, AverageMeter, MixUpWrapper, LR_Scheduler)
class Options():
def __init__(self):
# data settings
parser = argparse.ArgumentParser(description='Deep Encoding')
parser.add_argument('--dataset', type=str, default='cifar10',
help='training dataset (default: cifar10)')
parser.add_argument('--base-size', type=int, default=None,
help='base image size')
parser.add_argument('--crop-size', type=int, default=224,
help='crop image size')
parser.add_argument('--label-smoothing', type=float, default=0.0,
help='label-smoothing (default eta: 0.0)')
parser.add_argument('--mixup', type=float, default=0.0,
help='mixup (default eta: 0.0)')
parser.add_argument('--rand-aug', action='store_true',
default=False, help='random augment')
# model params
parser.add_argument('--model', type=str, default='densenet',
help='network model type (default: densenet)')
parser.add_argument('--rectify', action='store_true',
default=False, help='rectify convolution')
parser.add_argument('--rectify-avg', action='store_true',
default=False, help='rectify convolution')
parser.add_argument('--pretrained', action='store_true',
default=False, help='load pretrianed mode')
parser.add_argument('--last-gamma', action='store_true', default=False,
help='whether to init gamma of the last BN layer in \
each bottleneck to 0 (default: False)')
parser.add_argument('--dropblock-prob', type=float, default=0,
help='DropBlock prob. default is 0.')
parser.add_argument('--final-drop', type=float, default=0,
help='final dropout prob. default is 0.')
# training params
parser.add_argument('--batch-size', type=int, default=128, metavar='N',
help='batch size for training (default: 128)')
parser.add_argument('--test-batch-size', type=int, default=256, metavar='N',
help='batch size for testing (default: 256)')
parser.add_argument('--epochs', type=int, default=120, metavar='N',
help='number of epochs to train (default: 600)')
parser.add_argument('--start_epoch', type=int, default=0,
metavar='N', help='the epoch number to start (default: 1)')
parser.add_argument('--workers', type=int, default=8,
metavar='N', help='dataloader threads')
# optimizer
parser.add_argument('--lr', type=float, default=0.1, metavar='LR',
help='learning rate (default: 0.1)')
parser.add_argument('--lr-scheduler', type=str, default='cos',
help='learning rate scheduler (default: cos)')
parser.add_argument('--warmup-epochs', type=int, default=0,
help='number of warmup epochs (default: 0)')
parser.add_argument('--momentum', type=float, default=0.9,
metavar='M', help='SGD momentum (default: 0.9)')
parser.add_argument('--weight-decay', type=float, default=1e-4,
metavar ='M', help='SGD weight decay (default: 1e-4)')
parser.add_argument('--no-bn-wd', action='store_true',
default=False, help='no bias decay')
# seed
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
# checking point
parser.add_argument('--resume', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--checkname', type=str, default='default',
help='set the checkpoint name')
# distributed
parser.add_argument('--world-size', default=1, type=int,
help='number of nodes for distributed training')
parser.add_argument('--rank', default=0, type=int,
help='node rank for distributed training')
parser.add_argument('--dist-url', default='tcp://localhost:23456', type=str,
help='url used to set up distributed training')
parser.add_argument('--dist-backend', default='nccl', type=str,
help='distributed backend')
self.parser = parser
def parse(self):
args = self.parser.parse_args()
return args
def main():
args = Options().parse()
ngpus_per_node = torch.cuda.device_count()
args.world_size = ngpus_per_node * args.world_size
args.lr = args.lr * args.world_size
mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
# global variable
best_pred = 0.0
acclist_train = []
acclist_val = []
def main_worker(gpu, ngpus_per_node, args):
args.gpu = gpu
args.rank = args.rank * ngpus_per_node + gpu
print('rank: {} / {}'.format(args.rank, args.world_size))
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
torch.cuda.set_device(args.gpu)
# init the args
global best_pred, acclist_train, acclist_val
if args.gpu == 0:
print(args)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
# init dataloader
transform_train, transform_val = encoding.transforms.get_transform(
args.dataset, args.base_size, args.crop_size, args.rand_aug)
trainset = encoding.datasets.get_dataset(args.dataset, root=os.path.expanduser('~/.encoding/data'),
transform=transform_train, train=True, download=True)
valset = encoding.datasets.get_dataset(args.dataset, root=os.path.expanduser('~/.encoding/data'),
transform=transform_val, train=False, download=True)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
train_loader = torch.utils.data.DataLoader(
trainset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True,
sampler=train_sampler)
val_sampler = torch.utils.data.distributed.DistributedSampler(valset, shuffle=False)
val_loader = torch.utils.data.DataLoader(
valset, batch_size=args.test_batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True,
sampler=val_sampler)
# init the model
model_kwargs = {}
if args.pretrained:
model_kwargs['pretrained'] = True
if args.final_drop > 0.0:
model_kwargs['final_drop'] = args.final_drop
if args.dropblock_prob > 0.0:
model_kwargs['dropblock_prob'] = args.dropblock_prob
if args.last_gamma:
model_kwargs['last_gamma'] = True
if args.rectify:
model_kwargs['rectified_conv'] = True
model_kwargs['rectify_avg'] = args.rectify_avg
model = encoding.models.get_model(args.model, **model_kwargs)
if args.dropblock_prob > 0.0:
from functools import partial
from encoding.nn import reset_dropblock
nr_iters = (args.epochs - args.warmup_epochs) * len(train_loader)
apply_drop_prob = partial(reset_dropblock, args.warmup_epochs*len(train_loader),
nr_iters, 0.0, args.dropblock_prob)
model.apply(apply_drop_prob)
if args.gpu == 0:
print(model)
if args.mixup > 0:
train_loader = MixUpWrapper(args.mixup, 1000, train_loader, args.gpu)
criterion = NLLMultiLabelSmooth(args.label_smoothing)
elif args.label_smoothing > 0.0:
criterion = LabelSmoothing(args.label_smoothing)
else:
criterion = nn.CrossEntropyLoss()
model.cuda(args.gpu)
criterion.cuda(args.gpu)
model = DistributedDataParallel(model, device_ids=[args.gpu])
# criterion and optimizer
if args.no_bn_wd:
parameters = model.named_parameters()
param_dict = {}
for k, v in parameters:
param_dict[k] = v
bn_params = [v for n, v in param_dict.items() if ('bn' in n or 'bias' in n)]
rest_params = [v for n, v in param_dict.items() if not ('bn' in n or 'bias' in n)]
if args.gpu == 0:
print(" Weight decay NOT applied to BN parameters ")
print(f'len(parameters): {len(list(model.parameters()))} = {len(bn_params)} + {len(rest_params)}')
optimizer = torch.optim.SGD([{'params': bn_params, 'weight_decay': 0 },
{'params': rest_params, 'weight_decay': args.weight_decay}],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# check point
if args.resume is not None:
if os.path.isfile(args.resume):
if args.gpu == 0:
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] + 1 if args.start_epoch == 0 else args.start_epoch
best_pred = checkpoint['best_pred']
acclist_train = checkpoint['acclist_train']
acclist_val = checkpoint['acclist_val']
model.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
if args.gpu == 0:
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
raise RuntimeError ("=> no resume checkpoint found at '{}'".\
format(args.resume))
scheduler = LR_Scheduler(args.lr_scheduler,
base_lr=args.lr,
num_epochs=args.epochs,
iters_per_epoch=len(train_loader),
warmup_epochs=args.warmup_epochs)
def train(epoch):
train_sampler.set_epoch(epoch)
model.train()
losses = AverageMeter()
top1 = AverageMeter()
global best_pred, acclist_train
for batch_idx, (data, target) in enumerate(train_loader):
scheduler(optimizer, batch_idx, epoch, best_pred)
if not args.mixup:
data, target = data.cuda(args.gpu), target.cuda(args.gpu)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
if not args.mixup:
acc1 = accuracy(output, target, topk=(1,))
top1.update(acc1[0], data.size(0))
losses.update(loss.item(), data.size(0))
if batch_idx % 100 == 0 and args.gpu == 0:
if args.mixup:
print('Batch: %d| Loss: %.3f'%(batch_idx, losses.avg))
else:
print('Batch: %d| Loss: %.3f | Top1: %.3f'%(batch_idx, losses.avg, top1.avg))
acclist_train += [top1.avg]
def validate(epoch):
model.eval()
top1 = AverageMeter()
top5 = AverageMeter()
global best_pred, acclist_train, acclist_val
is_best = False
for batch_idx, (data, target) in enumerate(val_loader):
data, target = data.cuda(args.gpu), target.cuda(args.gpu)
with torch.no_grad():
output = model(data)
acc1, acc5 = accuracy(output, target, topk=(1, 5))
top1.update(acc1[0], data.size(0))
top5.update(acc5[0], data.size(0))
comm = MPI.COMM_WORLD
# send to master
sum1 = comm.gather(top1.sum, root=0)
cnt1 = comm.gather(top1.count, root=0)
sum5 = comm.gather(top5.sum, root=0)
cnt5 = comm.gather(top5.count, root=0)
# get back from master
sum1 = comm.bcast(sum1, root=0)
cnt1 = comm.bcast(cnt1, root=0)
sum5 = comm.bcast(sum5, root=0)
cnt5 = comm.bcast(cnt5, root=0)
if args.gpu == 0:
top1_acc = sum(sum1) / sum(cnt1)
top5_acc = sum(sum5) / len(cnt5)
print('Validation: Top1: %.3f | Top5: %.3f'%(top1_acc, top5_acc))
# save checkpoint
acclist_val += [top1_acc]
if top1_acc > best_pred:
best_pred = top1_acc
is_best = True
encoding.utils.save_checkpoint({
'epoch': epoch,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'best_pred': best_pred,
'acclist_train':acclist_train,
'acclist_val':acclist_val,
}, args=args, is_best=is_best)
for epoch in range(args.start_epoch, args.epochs):
tic = time.time()
train(epoch)
if epoch % 10 == 0:
validate(epoch)
elapsed = time.time() - tic
if args.gpu == 0:
print(f'Epoch: {epoch}, Time cost: {elapsed}')
validate(epoch)
if __name__ == "__main__":
main()
experiments/recognition/verify.py 0 → 100644
View file @ b872eb8c
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## Email: zhanghang0704@gmail.com
## Copyright (c) 2020
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
from __future__ import print_function
import os
import argparse
from tqdm import tqdm
import torch
import torch.nn as nn
import encoding
from encoding.utils import (accuracy, AverageMeter, MixUpWrapper, LR_Scheduler)
class Options():
def __init__(self):
# data settings
parser = argparse.ArgumentParser(description='Deep Encoding')
parser.add_argument('--dataset', type=str, default='cifar10',
help='training dataset (default: cifar10)')
parser.add_argument('--base-size', type=int, default=None,
help='base image size')
parser.add_argument('--crop-size', type=int, default=224,
help='crop image size')
# model params
parser.add_argument('--model', type=str, default='densenet',
help='network model type (default: densenet)')
parser.add_argument('--rectify', action='store_true',
default=False, help='rectify convolution')
parser.add_argument('--rectify-avg', action='store_true',
default=False, help='rectify convolution')
# training hyper params
parser.add_argument('--batch-size', type=int, default=128, metavar='N',
help='batch size for training (default: 128)')
parser.add_argument('--workers', type=int, default=32,
metavar='N', help='dataloader threads')
# cuda, seed and logging
parser.add_argument('--no-cuda', action='store_true',
default=False, help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
# checking point
parser.add_argument('--resume', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--verify', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--export', type=str, default=None,
help='put the path to resuming file if needed')
self.parser = parser
def parse(self):
args = self.parser.parse_args()
return args
def main():
# init the args
args = Options().parse()
args.cuda = not args.no_cuda and torch.cuda.is_available()
print(args)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
# init dataloader
_, transform_val = encoding.transforms.get_transform(args.dataset, args.base_size, args.crop_size)
valset = encoding.datasets.get_dataset(args.dataset, root=os.path.expanduser('~/.encoding/data'),
transform=transform_val, train=False, download=True)
val_loader = torch.utils.data.DataLoader(
valset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True if args.cuda else False)
# init the model
model_kwargs = {'pretrained': True}
if args.rectify:
model_kwargs['rectified_conv'] = True
model_kwargs['rectify_avg'] = args.rectify_avg
model = encoding.models.get_model(args.model, **model_kwargs)
print(model)
if args.cuda:
model.cuda()
# Please use CUDA_VISIBLE_DEVICES to control the number of gpus
model = nn.DataParallel(model)
# checkpoint
if args.verify:
if os.path.isfile(args.verify):
print("=> loading checkpoint '{}'".format(args.verify))
model.module.load_state_dict(torch.load(args.verify))
else:
raise RuntimeError ("=> no verify checkpoint found at '{}'".\
format(args.verify))
elif args.resume is not None:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model.module.load_state_dict(checkpoint['state_dict'])
else:
raise RuntimeError ("=> no resume checkpoint found at '{}'".\
format(args.resume))
model.eval()
top1 = AverageMeter()
top5 = AverageMeter()
is_best = False
tbar = tqdm(val_loader, desc='\r')
for batch_idx, (data, target) in enumerate(tbar):
if args.cuda:
data, target = data.cuda(), target.cuda()
with torch.no_grad():
output = model(data)
acc1, acc5 = accuracy(output, target, topk=(1, 5))
top1.update(acc1[0], data.size(0))
top5.update(acc5[0], data.size(0))
tbar.set_description('Top1: %.3f | Top5: %.3f'%(top1.avg, top5.avg))
print('Top1 Acc: %.3f | Top5 Acc: %.3f '%(top1.avg, top5.avg))
if args.export:
torch.save(model.module.state_dict(), args.export + '.pth')
if __name__ == "__main__":
main()
experiments/segmentation/test.py
View file @ b872eb8c
......@@ -5,6 +5,7 @@
###########################################################################
import os
import argparse
import numpy as np
from tqdm import tqdm
......@@ -16,10 +17,82 @@ from torch.nn.parallel.scatter_gather import gather
import encoding.utils as utils
from encoding.nn import SegmentationLosses, SyncBatchNorm
from encoding.parallel import DataParallelModel, DataParallelCriterion
from encoding.datasets import get_segmentation_dataset, test_batchify_fn
from encoding.datasets import get_dataset, test_batchify_fn
from encoding.models import get_model, get_segmentation_model, MultiEvalModule
from option import Options
class Options():
def __init__(self):
parser = argparse.ArgumentParser(description='PyTorch Segmentation')
# model and dataset
parser.add_argument('--model', type=str, default='encnet',
help='model name (default: encnet)')
parser.add_argument('--backbone', type=str, default='resnet50',
help='backbone name (default: resnet50)')
parser.add_argument('--dataset', type=str, default='ade20k',
help='dataset name (default: pascal12)')
parser.add_argument('--workers', type=int, default=16,
metavar='N', help='dataloader threads')
parser.add_argument('--base-size', type=int, default=520,
help='base image size')
parser.add_argument('--crop-size', type=int, default=480,
help='crop image size')
parser.add_argument('--train-split', type=str, default='train',
help='dataset train split (default: train)')
# training hyper params
parser.add_argument('--aux', action='store_true', default= False,
help='Auxilary Loss')
parser.add_argument('--se-loss', action='store_true', default= False,
help='Semantic Encoding Loss SE-loss')
parser.add_argument('--se-weight', type=float, default=0.2,
help='SE-loss weight (default: 0.2)')
parser.add_argument('--batch-size', type=int, default=16,
metavar='N', help='input batch size for \
training (default: auto)')
parser.add_argument('--test-batch-size', type=int, default=16,
metavar='N', help='input batch size for \
testing (default: same as batch size)')
# cuda, seed and logging
parser.add_argument('--no-cuda', action='store_true', default=
False, help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
# checking point
parser.add_argument('--resume', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--verify', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--model-zoo', type=str, default=None,
help='evaluating on model zoo model')
# evaluation option
parser.add_argument('--eval', action='store_true', default= False,
help='evaluating mIoU')
parser.add_argument('--export', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--acc-bn', action='store_true', default= False,
help='Re-accumulate BN statistics')
parser.add_argument('--test-val', action='store_true', default= False,
help='generate masks on val set')
parser.add_argument('--no-val', action='store_true', default= False,
help='skip validation during training')
# test option
parser.add_argument('--test-folder', type=str, default=None,
help='path to test image folder')
# the parser
self.parser = parser
def parse(self):
args = self.parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
print(args)
return args
@torch.no_grad()
def reset_bn_statistics(m):
if isinstance(m, torch.nn.BatchNorm2d):
#print(m)
m.momentum = 0.0
m.reset_running_stats()
def test(args):
# output folder
......@@ -32,14 +105,14 @@ def test(args):
transform.Normalize([.485, .456, .406], [.229, .224, .225])])
# dataset
if args.eval:
testset = get_segmentation_dataset(args.dataset, split='val', mode='testval',
transform=input_transform)
testset = get_dataset(args.dataset, split='val', mode='testval',
transform=input_transform)
elif args.test_val:
testset = get_segmentation_dataset(args.dataset, split='val', mode='test',
transform=input_transform)
testset = get_dataset(args.dataset, split='val', mode='test',
transform=input_transform)
else:
testset = get_segmentation_dataset(args.dataset, split='test', mode='test',
transform=input_transform)
testset = get_dataset(args.dataset, split='test', mode='test',
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
......@@ -53,20 +126,47 @@ def test(args):
#model.crop_size = args.crop_size
else:
model = get_segmentation_model(args.model, dataset=args.dataset,
backbone = args.backbone, aux = args.aux,
se_loss = args.se_loss, norm_layer = SyncBatchNorm,
backbone=args.backbone, aux = args.aux,
se_loss=args.se_loss,
norm_layer=torch.nn.BatchNorm2d if args.acc_bn else SyncBatchNorm,
base_size=args.base_size, crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
if args.verify is not None and os.path.isfile(args.verify):
print("=> loading checkpoint '{}'".format(args.verify))
model.load_state_dict(torch.load(args.verify))
elif args.resume is not None and os.path.isfile(args.resume):
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
else:
raise RuntimeError ("=> no checkpoint found")
print(model)
# accumulate bn statistics
if args.acc_bn:
print('Reseting BN statistics')
model.apply(reset_bn_statistics)
data_kwargs = {'transform': input_transform, 'base_size': args.base_size,
'crop_size': args.crop_size}
trainset = get_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs)
trainloader = data.DataLoader(trainset, batch_size=args.batch_size,
drop_last=True, shuffle=True, **loader_kwargs)
tbar = tqdm(trainloader)
model.train()
model.cuda()
for i, (image, dst) in enumerate(tbar):
image = image.cuda()
with torch.no_grad():
outputs = model(image)
if i > 1000: break
if args.export:
torch.save(model.state_dict(), args.export + '.pth')
return
scales = [0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0]
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75]#, 2.0
evaluator = MultiEvalModule(model, testset.num_class, scales=scales).cuda()
evaluator.eval()
metric = utils.SegmentationMetric(testset.num_class)
......@@ -89,6 +189,8 @@ def test(args):
outname = os.path.splitext(impath)[0] + '.png'
mask.save(os.path.join(outdir, outname))
print( 'pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU))
if __name__ == "__main__":
args = Options().parse()
torch.manual_seed(args.seed)
......
experiments/segmentation/train.py
View file @ b872eb8c
......@@ -6,6 +6,7 @@
import os
import copy
import argparse
import numpy as np
from tqdm import tqdm
......@@ -15,16 +16,113 @@ import torchvision.transforms as transform
from torch.nn.parallel.scatter_gather import gather
import encoding.utils as utils
from encoding.nn import SegmentationLosses, SyncBatchNorm, OHEMSegmentationLosses
from encoding.nn import SegmentationLosses, SyncBatchNorm
from encoding.parallel import DataParallelModel, DataParallelCriterion
from encoding.datasets import get_dataset
from encoding.models import get_segmentation_model
from option import Options
class Options():
def __init__(self):
parser = argparse.ArgumentParser(description='PyTorch \
Segmentation')
# model and dataset
parser.add_argument('--model', type=str, default='encnet',
help='model name (default: encnet)')
parser.add_argument('--backbone', type=str, default='resnet50',
help='backbone name (default: resnet50)')
parser.add_argument('--dataset', type=str, default='ade20k',
help='dataset name (default: pascal12)')
parser.add_argument('--workers', type=int, default=16,
metavar='N', help='dataloader threads')
parser.add_argument('--base-size', type=int, default=520,
help='base image size')
parser.add_argument('--crop-size', type=int, default=480,
help='crop image size')
parser.add_argument('--train-split', type=str, default='train',
help='dataset train split (default: train)')
# training hyper params
parser.add_argument('--aux', action='store_true', default= False,
help='Auxilary Loss')
parser.add_argument('--aux-weight', type=float, default=0.2,
help='Auxilary loss weight (default: 0.2)')
parser.add_argument('--se-loss', action='store_true', default= False,
help='Semantic Encoding Loss SE-loss')
parser.add_argument('--se-weight', type=float, default=0.2,
help='SE-loss weight (default: 0.2)')
parser.add_argument('--epochs', type=int, default=None, metavar='N',
help='number of epochs to train (default: auto)')
parser.add_argument('--start_epoch', type=int, default=0,
metavar='N', help='start epochs (default:0)')
parser.add_argument('--batch-size', type=int, default=16,
metavar='N', help='input batch size for \
training (default: auto)')
parser.add_argument('--test-batch-size', type=int, default=16,
metavar='N', help='input batch size for \
testing (default: same as batch size)')
# optimizer params
parser.add_argument('--lr', type=float, default=None, metavar='LR',
help='learning rate (default: auto)')
parser.add_argument('--lr-scheduler', type=str, default='poly',
help='learning rate scheduler (default: poly)')
parser.add_argument('--momentum', type=float, default=0.9,
metavar='M', help='momentum (default: 0.9)')
parser.add_argument('--weight-decay', type=float, default=1e-4,
metavar='M', help='w-decay (default: 1e-4)')
# cuda, seed and logging
parser.add_argument('--no-cuda', action='store_true', default=
False, help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
# checking point
parser.add_argument('--resume', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--checkname', type=str, default='default',
help='set the checkpoint name')
parser.add_argument('--model-zoo', type=str, default=None,
help='evaluating on model zoo model')
# finetuning pre-trained models
parser.add_argument('--ft', action='store_true', default= False,
help='finetuning on a different dataset')
# evaluation option
parser.add_argument('--eval', action='store_true', default= False,
help='evaluating mIoU')
parser.add_argument('--test-val', action='store_true', default= False,
help='generate masks on val set')
parser.add_argument('--no-val', action='store_true', default= False,
help='skip validation during training')
# test option
parser.add_argument('--test-folder', type=str, default=None,
help='path to test image folder')
# the parser
self.parser = parser
def parse(self):
args = self.parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
# default settings for epochs, batch_size and lr
if args.epochs is None:
epoches = {
'coco': 30,
'pascal_aug': 80,
'pascal_voc': 50,
'pcontext': 80,
'ade20k': 180,
'citys': 240,
}
args.epochs = epoches[args.dataset.lower()]
if args.lr is None:
lrs = {
'coco': 0.004,
'pascal_aug': 0.001,
'pascal_voc': 0.0001,
'pcontext': 0.001,
'ade20k': 0.004,
'citys': 0.004,
}
args.lr = lrs[args.dataset.lower()] / 16 * args.batch_size
print(args)
return args
torch_ver = torch.__version__[:3]
if torch_ver == '0.3':
from torch.autograd import Variable
class Trainer():
def __init__(self, args):
......@@ -36,10 +134,8 @@ class Trainer():
# dataset
data_kwargs = {'transform': input_transform, 'base_size': args.base_size,
'crop_size': args.crop_size}
trainset = get_dataset(args.dataset, split=args.train_split, mode='train',
**data_kwargs)
testset = get_dataset(args.dataset, split='val', mode ='val',
**data_kwargs)
trainset = get_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs)
testset = get_dataset(args.dataset, split='val', mode ='val', **data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
......@@ -92,8 +188,8 @@ class Trainer():
if args.ft:
args.start_epoch = 0
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.trainloader))
self.scheduler = utils.LR_Scheduler_Head(args.lr_scheduler, args.lr,
args.epochs, len(self.trainloader))
self.best_pred = 0.0
def training(self, epoch):
......@@ -103,9 +199,6 @@ class Trainer():
for i, (image, target) in enumerate(tbar):
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
if torch_ver == "0.3":
image = Variable(image)
target = Variable(target)
outputs = self.model(image)
loss = self.criterion(outputs, target)
loss.backward()
......@@ -140,12 +233,8 @@ class Trainer():
total_inter, total_union, total_correct, total_label = 0, 0, 0, 0
tbar = tqdm(self.valloader, desc='\r')
for i, (image, target) in enumerate(tbar):
if torch_ver == "0.3":
image = Variable(image, volatile=True)
with torch.no_grad():
correct, labeled, inter, union = eval_batch(self.model, image, target)
else:
with torch.no_grad():
correct, labeled, inter, union = eval_batch(self.model, image, target)
total_correct += correct
total_label += labeled
......
experiments/segmentation/option.py → experiments/segmentation/train_dist.py
View file @ b872eb8c
###########################################################################
# Created by: Hang Zhang
# Email: zhang.hang@rutgers.edu
# Copyright (c) 2017
###########################################################################
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## Email: zhanghang0704@gmail.com
## Copyright (c) 2020
##
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import os
import time
import argparse
import numpy as np
from tqdm import tqdm
#from mpi4py import MPI
import torch
from torch.utils import data
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.backends.cudnn as cudnn
import torchvision.transforms as transform
from torch.nn.parallel.scatter_gather import gather
from torch.nn.parallel import DistributedDataParallel
import encoding.utils as utils
from encoding.nn import SegmentationLosses, DistSyncBatchNorm
from encoding.datasets import get_dataset
from encoding.models import get_segmentation_model
class Options():
def __init__(self):
......@@ -17,13 +38,13 @@ class Options():
help='model name (default: encnet)')
parser.add_argument('--backbone', type=str, default='resnet50',
help='backbone name (default: resnet50)')
parser.add_argument('--rectify', action='store_true',
default=False, help='rectify convolution')
parser.add_argument('--rectify-avg', action='store_true',
default=False, help='rectify convolution')
parser.add_argument('--dataset', type=str, default='ade20k',
help='dataset name (default: pascal12)')
parser.add_argument('--data-folder', type=str,
default=os.path.join(os.environ['HOME'], 'data'),
help='training dataset folder (default: \
$(HOME)/data)')
parser.add_argument('--workers', type=int, default=16,
parser.add_argument('--workers', type=int, default=8,
metavar='N', help='dataloader threads')
parser.add_argument('--base-size', type=int, default=520,
help='base image size')
......@@ -44,7 +65,7 @@ class Options():
help='number of epochs to train (default: auto)')
parser.add_argument('--start_epoch', type=int, default=0,
metavar='N', help='start epochs (default:0)')
parser.add_argument('--batch-size', type=int, default=16,
parser.add_argument('--batch-size', type=int, default=2,
metavar='N', help='input batch size for \
training (default: auto)')
parser.add_argument('--test-batch-size', type=int, default=16,
......@@ -60,8 +81,6 @@ class Options():
parser.add_argument('--weight-decay', type=float, default=1e-4,
metavar='M', help='w-decay (default: 1e-4)')
# cuda, seed and logging
parser.add_argument('--no-cuda', action='store_true', default=
False, help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
# checking point
......@@ -79,17 +98,23 @@ class Options():
help='evaluating mIoU')
parser.add_argument('--test-val', action='store_true', default= False,
help='generate masks on val set')
parser.add_argument('--no-val', action='store_true', default= False,
help='skip validation during training')
# test option
parser.add_argument('--test-folder', type=str, default=None,
help='path to test image folder')
# distributed
parser.add_argument('--world-size', default=1, type=int,
help='number of nodes for distributed training')
parser.add_argument('--rank', default=0, type=int,
help='node rank for distributed training')
parser.add_argument('--dist-url', default='tcp://localhost:23456', type=str,
help='url used to set up distributed training')
parser.add_argument('--dist-backend', default='nccl', type=str,
help='distributed backend')
# the parser
self.parser = parser
def parse(self):
args = self.parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
# default settings for epochs, batch_size and lr
if args.epochs is None:
epoches = {
......@@ -97,7 +122,7 @@ class Options():
'pascal_aug': 80,
'pascal_voc': 50,
'pcontext': 80,
'ade20k': 180,
'ade20k': 120,
'citys': 240,
}
args.epochs = epoches[args.dataset.lower()]
......@@ -107,9 +132,195 @@ class Options():
'pascal_aug': 0.001,
'pascal_voc': 0.0001,
'pcontext': 0.001,
'ade20k': 0.004,
'citys': 0.004,
'ade20k': 0.01,
'citys': 0.01,
}
args.lr = lrs[args.dataset.lower()] / 16 * args.batch_size
print(args)
return args
#def mpi_avg_all(*args):
# comm = MPI.COMM_WORLD
# # send to master
# sum_args = []
# for arg in args:
# sum_args.append(sum(comm.gather(arg, root=0)))
# sum_args = [item / len(args) for item in sum_args]
# return tuple(sum_args)
def torch_dist_avg(*args):
process_group = torch.distributed.group.WORLD
tensor_args = []
pending_res = []
for arg in args:
tensor_arg = torch.tensor(arg)
tensor_args.append(tensor_arg)
pending_res.append(torch.distributed.all_reduce(tensor_arg, group=process_group, async_op=True))
for res in pending_res:
res.wait()
ret = [x.item()/len(tensor_args) for x in tensor_args]
return ret
def main():
args = Options().parse()
ngpus_per_node = torch.cuda.device_count()
args.world_size = ngpus_per_node * args.world_size
args.lr = args.lr * args.world_size
mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
best_pred = 0.0
def main_worker(gpu, ngpus_per_node, args):
global best_pred
args.gpu = gpu
args.rank = args.rank * ngpus_per_node + gpu
print('rank: {} / {}'.format(args.rank, args.world_size))
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
torch.cuda.set_device(args.gpu)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
cudnn.benchmark = True
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])])
# dataset
data_kwargs = {'transform': input_transform, 'base_size': args.base_size,
'crop_size': args.crop_size}
trainset = get_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs)
valset = get_dataset(args.dataset, split='val', mode ='val', **data_kwargs)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
val_sampler = torch.utils.data.distributed.DistributedSampler(valset, shuffle=False)
# dataloader
loader_kwargs = {'batch_size': args.batch_size, 'num_workers': args.workers, 'pin_memory': True}
trainloader = data.DataLoader(trainset, sampler=train_sampler, drop_last=True, **loader_kwargs)
valloader = data.DataLoader(valset, sampler=val_sampler, **loader_kwargs)
nclass = trainset.num_class
# model
model_kwargs = {}
if args.rectify:
model_kwargs['rectified_conv'] = True
model_kwargs['rectify_avg'] = args.rectify_avg
model = get_segmentation_model(args.model, dataset=args.dataset,
backbone=args.backbone, aux=args.aux,
se_loss=args.se_loss, norm_layer=DistSyncBatchNorm,
base_size=args.base_size, crop_size=args.crop_size,
**model_kwargs)
if args.gpu == 0:
print(model)
# optimizer using different LR
params_list = [{'params': model.pretrained.parameters(), 'lr': args.lr},]
if hasattr(model, 'head'):
params_list.append({'params': model.head.parameters(), 'lr': args.lr*10})
if hasattr(model, 'auxlayer'):
params_list.append({'params': model.auxlayer.parameters(), 'lr': args.lr*10})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# criterions
criterion = SegmentationLosses(se_loss=args.se_loss,
aux=args.aux,
nclass=nclass,
se_weight=args.se_weight,
aux_weight=args.aux_weight)
# distributed data parallel
model.cuda(args.gpu)
criterion.cuda(args.gpu)
model = DistributedDataParallel(model, device_ids=[args.gpu])
metric = utils.SegmentationMetric(nclass=nclass)
# resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
if not args.ft:
optimizer.load_state_dict(checkpoint['optimizer'])
best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
# clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
# lr scheduler
scheduler = utils.LR_Scheduler_Head(args.lr_scheduler, args.lr,
args.epochs, len(trainloader))
def training(epoch):
global best_pred
train_loss = 0.0
model.train()
tic = time.time()
for i, (image, target) in enumerate(trainloader):
scheduler(optimizer, i, epoch, best_pred)
optimizer.zero_grad()
outputs = model(image)
target = target.cuda(args.gpu)
loss = criterion(*outputs, target)
loss.backward()
optimizer.step()
train_loss += loss.item()
if i % 100 == 0 and args.gpu == 0:
iter_per_sec = 100.0 / (time.time() - tic) if i != 0 else 1.0/ (time.time() - tic)
tic = time.time()
print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Train loss: {:.3f}'. \
format(epoch, i, iter_per_sec, train_loss / (i + 1)))
def validation(epoch):
# Fast test during the training using single-crop only
global best_pred
is_best = False
model.eval()
metric.reset()
for i, (image, target) in enumerate(valloader):
with torch.no_grad():
#correct, labeled, inter, union = eval_batch(model, image, target)
pred = model(image)[0]
target = target.cuda(args.gpu)
metric.update(target, pred)
pixAcc, mIoU = metric.get()
if i % 100 == 0 and args.gpu == 0:
print('pixAcc: %.3f, mIoU: %.3f' % (pixAcc, mIoU))
if args.gpu == 0:
pixAcc, mIoU = torch_dist_avg(pixAcc, mIoU)
print('pixAcc: %.3f, mIoU: %.3f' % (pixAcc, mIoU))
new_pred = (pixAcc + mIoU)/2
if new_pred > best_pred:
is_best = True
best_pred = new_pred
utils.save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'best_pred': best_pred,
}, args, is_best)
if args.gpu == 0:
print('Starting Epoch:', args.start_epoch)
print('Total Epoches:', args.epochs)
for epoch in range(args.start_epoch, args.epochs):
tic = time.time()
training(epoch)
if epoch % 10 == 0:
validation(epoch)
elapsed = time.time() - tic
if args.gpu == 0:
print(f'Epoch: {epoch}, Time cost: {elapsed}')
validation(epoch)
if __name__ == "__main__":
main()
scripts/prepare_citys.py 0 → 100644
View file @ b872eb8c
"""Prepare Cityscapes dataset"""
import os
import shutil
import argparse
import zipfile
from encoding.utils import check_sha1, download, mkdir
_TARGET_DIR = os.path.expanduser('~/.encoding/data')
def parse_args():
parser = argparse.ArgumentParser(
description='Initialize ADE20K dataset.',
epilog='Example: python prepare_cityscapes.py',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--download-dir', default=None, help='dataset directory on disk')
args = parser.parse_args()
return args
def download_city(path, overwrite=False):
_CITY_DOWNLOAD_URLS = [
('gtFine_trainvaltest.zip', '99f532cb1af174f5fcc4c5bc8feea8c66246ddbc'),
('leftImg8bit_trainvaltest.zip', '2c0b77ce9933cc635adda307fbba5566f5d9d404')]
download_dir = os.path.join(path, 'downloads')
mkdir(download_dir)
for filename, checksum in _CITY_DOWNLOAD_URLS:
if not check_sha1(filename, checksum):
raise UserWarning('File {} is downloaded but the content hash does not match. ' \
'The repo may be outdated or download may be incomplete. ' \
'If the "repo_url" is overridden, consider switching to ' \
'the default repo.'.format(filename))
# extract
with zipfile.ZipFile(filename,"r") as zip_ref:
zip_ref.extractall(path=path)
print("Extracted", filename)
if __name__ == '__main__':
args = parse_args()
mkdir(os.path.expanduser('~/.encoding/data'))
if args.download_dir is not None:
if os.path.isdir(_TARGET_DIR):
os.remove(_TARGET_DIR)
# make symlink
os.symlink(args.download_dir, _TARGET_DIR)
else:
download_city(_TARGET_DIR, overwrite=False)
scripts/prepare_imagenet.py
View file @ b872eb8c
......@@ -6,9 +6,10 @@ import pickle
import gzip
import subprocess
from tqdm import tqdm
import subprocess
from encoding.utils import check_sha1, download, mkdir
_TARGET_DIR = os.path.expanduser('~/.encoding/datasets/imagenet')
_TARGET_DIR = os.path.expanduser('~/.encoding/data/ILSVRC2012')
_TRAIN_TAR = 'ILSVRC2012_img_train.tar'
_TRAIN_TAR_SHA1 = '43eda4fe35c1705d6606a6a7a633bc965d194284'
_VAL_TAR = 'ILSVRC2012_img_val.tar'
......@@ -37,42 +38,8 @@ def check_file(filename, checksum, sha1):
if checksum and not check_sha1(filename, sha1):
raise ValueError('Corrupted file: '+filename)
def build_rec_process(img_dir, train=False, num_thread=1):
rec_dir = os.path.abspath(os.path.join(img_dir, '../rec'))
mkdir(rec_dir)
prefix = 'train' if train else 'val'
print('Building ImageRecord file for ' + prefix + ' ...')
to_path = rec_dir
# download lst file and im2rec script
script_path = os.path.join(rec_dir, 'im2rec.py')
script_url = 'https://raw.githubusercontent.com/apache/incubator-encoding/master/tools/im2rec.py'
download(script_url, script_path)
lst_path = os.path.join(rec_dir, prefix + '.lst')
lst_url = 'http://data.encoding.io/models/imagenet/resnet/' + prefix + '.lst'
download(lst_url, lst_path)
# execution
import sys
cmd = [
sys.executable,
script_path,
rec_dir,
img_dir,
'--recursive',
'--pass-through',
'--pack-label',
'--num-thread',
str(num_thread)
]
subprocess.call(cmd)
os.remove(script_path)
os.remove(lst_path)
print('ImageRecord file for ' + prefix + ' has been built!')
def extract_train(tar_fname, target_dir, with_rec=False, num_thread=1):
os.makedirs(target_dir)
mkdir(target_dir)
with tarfile.open(tar_fname) as tar:
print("Extracting "+tar_fname+"...")
# extract each class one-by-one
......@@ -88,32 +55,24 @@ def extract_train(tar_fname, target_dir, with_rec=False, num_thread=1):
os.remove(class_fname)
pbar.update(1)
pbar.close()
if with_rec:
build_rec_process(target_dir, True, num_thread)
def extract_val(tar_fname, target_dir, with_rec=False, num_thread=1):
os.makedirs(target_dir)
mkdir(target_dir)
print('Extracting ' + tar_fname)
with tarfile.open(tar_fname) as tar:
tar.extractall(target_dir)
# build rec file before images are moved into subfolders
if with_rec:
build_rec_process(target_dir, False, num_thread)
# move images to proper subfolders
val_maps_file = os.path.join(os.path.dirname(__file__), 'imagenet_val_maps.pklz')
with gzip.open(val_maps_file, 'rb') as f:
dirs, mappings = pickle.load(f)
for d in dirs:
os.makedirs(os.path.join(target_dir, d))
for m in mappings:
os.rename(os.path.join(target_dir, m[0]), os.path.join(target_dir, m[1], m[0]))
subprocess.call(["wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash"],
cwd=target_dir, shell=True)
def main():
args = parse_args()
target_dir = os.path.expanduser(args.target_dir)
if os.path.exists(target_dir):
raise ValueError('Target dir ['+target_dir+'] exists. Remove it first')
#if os.path.exists(target_dir):
# raise ValueError('Target dir ['+target_dir+'] exists. Remove it first')
download_dir = os.path.expanduser(args.download_dir)
train_tar_fname = os.path.join(download_dir, _TRAIN_TAR)
......
setup.py
View file @ b872eb8c
......@@ -18,7 +18,7 @@ import setuptools.command.install
cwd = os.path.dirname(os.path.abspath(__file__))
version = '1.1.2'
version = '1.2.0'
try:
from datetime import date
today = date.today()
......@@ -52,14 +52,7 @@ requirements = [
'numpy',
'tqdm',
'nose',
'torch>=0.5.0',
'cffi>=1.0.0',
]
requirements = [
'numpy',
'tqdm',
'nose',
'portalocker',
'torch>=1.4.0',
'torchvision>=0.5.0',
'Pillow',
......
tests/unit_test/test_dataset.py 0 → 100644
View file @ b872eb8c
import argparse
from tqdm import tqdm
from torch.utils import data
import torchvision.transforms as transform
from encoding.datasets import get_segmentation_dataset
def main():
parser = argparse.ArgumentParser(description='Test Dataset.')
parser.add_argument('--dataset', type=str, default='ade20k',
help='dataset name (default: pascal12)')
args = parser.parse_args()
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])])
trainset = get_segmentation_dataset(args.dataset, split='val', mode='train',
transform=input_transform)
trainloader = data.DataLoader(trainset, batch_size=16,
drop_last=True, shuffle=True)
tbar = tqdm(trainloader)
max_label = -10
for i, (image, target) in enumerate(tbar):
tmax = target.max().item()
tmin = target.min().item()
assert(tmin >= -1)
if tmax > max_label:
max_label = tmax
assert(max_label < trainset.NUM_CLASS)
tbar.set_description("Batch %d, max label %d"%(i, max_label))
if __name__ == "__main__":
main()
tests/unit_test/test_function.py
View file @ b872eb8c
......@@ -46,132 +46,6 @@ def test_scaled_l2():
test = gradcheck(encoding.functions.scaled_l2, input, eps=EPS, atol=ATOL)
print('Testing scaled_l2(): {}'.format(test))
def test_aggregate_v2():
def py_aggregate_v2(A, X, C, STD, S):
B, N, D = X.size()
K = C.size(0)
#e_{k} = \sum_{i=1}^{N} a_{ik} (x_i - d_k) / \sigma_k
R = (X.view(B, N, 1, D).expand(B, N, K, D) - \
C.view(1, 1, K, D).expand(B, N, K, D)) / STD.view(1, 1, K, D)
#E = 1.0 / torch.sqrt(S + 1e-5).unsqueeze(0).unsqueeze(2) * (A.unsqueeze(3) * R).sum(1)
E2 = (A.unsqueeze(3) * R).sum(1)
return E2
B,N,K,D = 2,3,4,5
A = Variable(torch.cuda.DoubleTensor(B,N,K).uniform_(-0.5,0.5),
requires_grad=True)
X = Variable(torch.cuda.DoubleTensor(B,N,D).uniform_(-0.5,0.5),
requires_grad=True)
C = Variable(torch.cuda.DoubleTensor(K,D).uniform_(-0.5,0.5),
requires_grad=True)
STD = Variable(torch.cuda.DoubleTensor(K,D).uniform_(-0.5,0.5),
requires_grad=True)
S = Variable(torch.cuda.DoubleTensor(K).uniform_(-0.5,0.5),
requires_grad=True)
A2 = torch.from_numpy(A.detach().cpu().numpy()).cuda()
X2 = torch.from_numpy(X.detach().cpu().numpy()).cuda()
C2 = torch.from_numpy(C.detach().cpu().numpy()).cuda()
STD2 = torch.from_numpy(STD.detach().cpu().numpy()).cuda()
S2 = torch.from_numpy(S.detach().cpu().numpy()).cuda()
A2.requires_grad_()
X2.requires_grad_()
C2.requires_grad_()
STD2.requires_grad_()
S2.requires_grad_()
E = encoding.functions.aggregate_v2(A, X, C, STD)
E2 = py_aggregate_v2(A2, X2, C2, STD2, S2)
_assert_tensor_close(E.detach(), E2.detach())
input = (A, X, C, STD)
test = gradcheck(encoding.functions.aggregate_v2, input, eps=EPS, atol=ATOL)
print('Testing aggregate_v2(): {}'.format(test))
def test_encoding_dist():
def mahalanobis_dist(X, C):
B, N, D = X.size()
K = C.size(0)
# X \in BxNxD, C \in KxD
R = X.view(B, N, 1, D).expand(B, N, K, D) - \
C.view(1, 1, K, D).expand(B, N, K, D)
STD = torch.sqrt(R.pow(2).mean(0).mean(0) + 1e-6)
KD = (R / STD.view(1,1,K,D)).pow(2).sum(3)
return KD, STD
B,N,K,D = 2,3,4,5
RVar = torch.cuda.DoubleTensor(K,D).zero_()
X = torch.cuda.DoubleTensor(B,N,D).uniform_(-0.5,0.5)
C = torch.cuda.DoubleTensor(K,D).uniform_(-0.5,0.5)
X.requires_grad_()
C.requires_grad_()
X2 = torch.from_numpy(X.detach().cpu().numpy()).cuda()
C2 = torch.from_numpy(C.detach().cpu().numpy()).cuda()
X2.requires_grad_()
C2.requires_grad_()
# assert numeric correctness
KD, STD, Var_ = encoding.functions.encoding_dist(X, C, 1e-6)
KD2, STD2 = mahalanobis_dist(X2, C2)
_assert_tensor_close(STD.detach(), STD2.detach())
_assert_tensor_close(KD.detach(), KD2.detach())
# check backward
loss1 = KD.pow(2).sum() + STD.sum()
loss1.backward()
loss2 = KD2.pow(2).sum() + STD2.sum()
loss2.backward()
_assert_tensor_close(X.grad.detach(), X2.grad.detach())
_assert_tensor_close(C.grad.detach(), C2.grad.detach())
input = (X, C, 1e-6)
test = gradcheck(encoding.functions.encoding_dist, input, eps=EPS, atol=ATOL)
print('Testing encoding_dist(): {}'.format(test))
def test_encoding_dist_inference():
def mahalanobis_dist(X, C, STD):
B, N, D = X.size()
K = C.size(0)
# X \in BxNxD, C \in KxD
R = X.view(B, N, 1, D).expand(B, N, K, D) - \
C.view(1, 1, K, D).expand(B, N, K, D)
#STD = torch.sqrt(R.pow(2).mean(0).mean(0) + 1e-6)
KD = (R / STD.view(1,1,K,D)).pow(2).sum(3)
return KD
B,N,K,D = 2,3,4,5
X = Variable(torch.cuda.DoubleTensor(B,N,D).uniform_(-0.5,0.5),
requires_grad=True)
C = Variable(torch.cuda.DoubleTensor(K,D).uniform_(-0.5,0.5),
requires_grad=True)
STD = Variable(torch.cuda.DoubleTensor(K,D).uniform_(-0.5,0.5),
requires_grad=True)
X2 = torch.from_numpy(X.detach().cpu().numpy()).cuda()
C2 = torch.from_numpy(C.detach().cpu().numpy()).cuda()
STD2 = torch.from_numpy(STD.detach().cpu().numpy()).cuda()
X2.requires_grad_()
C2.requires_grad_()
STD2.requires_grad_()
E = encoding.functions.encoding_dist_inference(X, C, STD)
E2 = mahalanobis_dist(X2, C2, STD2)
loss1 = E.pow(2).sum()
loss2 = E2.pow(2).sum()
loss1.backward()
loss2.backward()
print('X.grad', X.grad)
print('X2.grad', X2.grad)
_assert_tensor_close(E.detach(), E2.detach())
_assert_tensor_close(X.grad.detach(), X2.grad.detach())
_assert_tensor_close(C.grad.detach(), C2.grad.detach())
_assert_tensor_close(STD.grad.detach(), STD2.grad.detach())
input = (X, C, STD)
test = gradcheck(encoding.functions.encoding_dist_inference, input, eps=EPS, atol=ATOL)
print('Testing encoding_dist_inference(): {}'.format(test))
def test_moments():
B,C,H = 2,3,4
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment