def __init__(self, cin, cout, sobel, net_heads=None, pool_size=None):
net_heads = net_heads if net_heads is not None else cfg.net_heads
pool_size = pool_size if pool_size is not None else cfg.net_avgpool_size
logger.debug('Backbone will be created wit the following heads: %s' %
net_heads)
# do init
super(ResNet34, self).__init__()
# build sobel
self.sobel = self._make_sobel_() if sobel else None
# build trunk net
self.inplanes = 64
self.layer1 = nn.Sequential(
nn.Conv2d(2 if sobel else cin,
64,
kernel_size=3,
stride=1,
padding=1,
bias=False), nn.BatchNorm2d(64,
track_running_stats=True),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2, padding=1))
self.layer2 = self._make_layer(BasicBlock, 64, 3)
self.layer3 = self._make_layer(BasicBlock, 128, 4, stride=2)
self.layer4 = self._make_layer(BasicBlock, 256, 6, stride=2)
self.layer5 = self._make_layer(BasicBlock, 512, 3, stride=2)
self.avgpool = nn.Sequential(nn.AvgPool2d(pool_size, stride=1),
Flatten())
heads = [
nn.Sequential(nn.Linear(512 * BasicBlock.expansion, head),
nn.Softmax(dim=1)) for head in net_heads
]
self.heads = nn.ModuleList(heads)
def __init__(self, cin, cout, sobel, grayscale, net_heads=None):
net_heads = net_heads if net_heads is not None else cfg.net_heads
logger.debug('Backbone will be created wit the following heads: %s' %
net_heads)
# do init
super(ResNet50Large, self).__init__()
# build sobel
self.sobel = self._make_sobel_() if sobel else None
self.grayscale = self._make_grayscale_() if grayscale else None
# build trunk net
self.inplanes = 64
if sobel:
input_channels = 2
elif grayscale:
input_channels = 1
else:
input_channels = cin
self.layer1 = nn.Sequential(
nn.Conv2d(input_channels,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False), nn.BatchNorm2d(64,
track_running_stats=True),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
self.layer2 = self._make_layer(Bottleneck, 64, 3)
self.layer3 = self._make_layer(Bottleneck, 128, 4, stride=2)
self.layer4 = self._make_layer(Bottleneck, 256, 6, stride=2)
self.layer5 = self._make_layer(Bottleneck, 512, 3, stride=2)
self.layer6 = self._make_layer(Bottleneck, 512, 3, stride=2)
self.avgpool = nn.Sequential(nn.AdaptiveAvgPool2d(1), Flatten())
heads = [
nn.Sequential(nn.Linear(512 * Bottleneck.expansion, head),
nn.Softmax(dim=1)) for head in net_heads
]
self.heads = nn.ModuleList(heads)
def main():
logger.info('Start to declare training variable')
cfg.device = 'cuda' if torch.cuda.is_available() else 'cpu'
logger.info('Session will be ran in device: [%s]' % cfg.device)
start_epoch = 0
best_acc = 0.
logger.info('Start to prepare data')
# get transformers
# train_transform is for data perturbation
train_transform = transforms.get(train=True)
# test_transform is for evaluation
test_transform = transforms.get(train=False)
# reduced_transform is for original training data
reduced_transform = get_reduced_transform(cfg.tfm_resize, cfg.tfm_size,
cfg.tfm_means, cfg.tfm_stds)
# get datasets
# each head should have its own trainset
train_splits = dict(cifar100=[['train', 'test']],
cifar10=[['train', 'test']],
stl10=[['train+unlabeled', 'test'], ['train', 'test']])
test_splits = dict(cifar100=['train', 'test'],
cifar10=['train', 'test'],
stl10=['train', 'test'])
# instance dataset for each head
if cfg.dataset.startswith('stl') or cfg.dataset.startswith('cifar'):
# otrainset: original trainset
otrainset = [
ConcatDataset([
datasets.get(split=split, transform=reduced_transform)
for split in train_splits[cfg.dataset][hidx]
]) for hidx in xrange(len(train_splits[cfg.dataset]))
]
# ptrainset: perturbed trainset
ptrainset = [
ConcatDataset([
datasets.get(split=split, transform=train_transform)
for split in train_splits[cfg.dataset][hidx]
]) for hidx in xrange(len(train_splits[cfg.dataset]))
]
# testset
testset = ConcatDataset([
datasets.get(split=split, transform=test_transform)
for split in test_splits[cfg.dataset]
])
else:
otrainset = [
ImageFolder(root=cfg.data_root, transform=reduced_transform)
for hidx in xrange(2)
]
ptrainset = [
ImageFolder(root=cfg.data_root, transform=train_transform)
for hidx in xrange(2)
]
testset = ImageFolder(root=cfg.data_root, transform=test_transform)
logger.debug(
'Dataset [%s] from directory [%s] is declared and %d samples '
'are loaded' % (cfg.dataset, cfg.data_root, len(testset)))
# declare data loaders for testset only
test_loader = DataLoader(testset,
batch_size=cfg.batch_size,
shuffle=False,
num_workers=cfg.num_workers)
logger.info('Start to build model')
net = networks.get()
criterion = DCLoss(cfg.dc_lamda)
optimizer = optimizers.get(
params=[val for _, val in net.trainable_parameters().iteritems()])
lr_handler = lr_policy.get()
# load session if checkpoint is provided
if cfg.resume:
assert os.path.exists(cfg.resume), "Resume file not found"
ckpt = torch.load(cfg.resume)
logger.info('Start to resume session for file: [%s]' % cfg.resume)
net.load_state_dict(ckpt['net'])
best_acc = ckpt['acc']
start_epoch = ckpt['epoch']
# data parallel
if cfg.device == 'cuda' and len(cfg.gpus.split(',')) > 1:
logger.info('Data parallel will be used for acceleration purpose')
device_ids = range(len(cfg.gpus.split(',')))
if not (hasattr(net, 'data_parallel')
and net.data_parallel(device_ids)):
net = nn.DataParallel(net, device_ids=device_ids)
cudnn.benchmark = True
else:
logger.info('Data parallel will not be used for acceleration')
# move modules to target device
net, criterion = net.to(cfg.device), criterion.to(cfg.device)
# tensorboard wrtier
writer = SummaryWriter(cfg.debug, log_dir=cfg.tfb_dir)
# start training
lr = cfg.base_lr
epoch = start_epoch
while lr > 0 and epoch < cfg.max_epochs:
lr = lr_handler.update(epoch, optimizer)
writer.add_scalar('Train/Learing_Rate', lr, epoch)
logger.info('Start to train at %d epoch with learning rate %.5f' %
(epoch, lr))
train(epoch, net, otrainset, ptrainset, optimizer, criterion, writer)
logger.info('Start to evaluate after %d epoch of training' % epoch)
acc, nmi, ari = evaluate(net, test_loader)
logger.info('Evaluation results at epoch %d are: '
'ACC: %.3f, NMI: %.3f, ARI: %.3f' % (epoch, acc, nmi, ari))
writer.add_scalar('Evaluate/ACC', acc, epoch)
writer.add_scalar('Evaluate/NMI', nmi, epoch)
writer.add_scalar('Evaluate/ARI', ari, epoch)
epoch += 1
if cfg.debug:
continue
# save checkpoint
is_best = acc > best_acc
best_acc = max(best_acc, acc)
if is_best:
best_epoch = epoch
save_checkpoint(
{
'net': net.state_dict(),
'optimizer': optimizer.state_dict(),
'acc': acc,
'epoch': epoch
},
is_best=is_best)
logger.info('Best evaluation results at epoch %d are: ACC: %.3f' %
(best_epoch, best_acc))
logger.info('Done')