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')