def test(args): # output folder outdir = 'outdir' if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), 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) else: testset = get_segmentation_dataset(args.dataset, split='test', mode='test', transform=input_transform) # dataloader loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm2d, 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']) utils.save_checkpoint({ 'state_dict': checkpoint['state_dict'], }, self.args, is_best, 'DANet101_reduce.pth.tar') print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch']))
def test(args): # data transforms input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, dilated=args.dilated, lateral=args.lateral, jpu=args.jpu, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm, 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'], strict=False) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) print(model) scales = [0.5, 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] if not args.ms: scales = [1.0] num_classes = datasets[args.dataset.lower()].NUM_CLASS evaluator = MultiEvalModule(model, num_classes, scales=scales, flip=args.ms).cuda() evaluator.eval() img = input_transform(Image.open( args.input_path).convert('RGB')).unsqueeze(0) with torch.no_grad(): output = evaluator.parallel_forward(img)[0] predict = torch.max(output, 1)[1].cpu().numpy() mask = utils.get_mask_pallete(predict, args.dataset) mask.save(args.save_path)
def __init__(self, args): self.args = args args.log_name = str(args.checkname) root_dir = getattr(args, "data_root", '../datasets') wo_head = getattr(args, "resume_wo_head", False) self.logger = utils.create_logger(args.log_root, args.log_name) # 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, 'logger': self.logger, 'scale': args.scale } trainset = get_segmentation_dataset(args.dataset, split='train', mode='train', root=root_dir, **data_kwargs) testset = get_segmentation_dataset(args.dataset, split='val', mode='val', root=root_dir, **data_kwargs) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, **kwargs) self.valloader = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) self.nclass = trainset.num_class # model model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm2d, base_size=args.base_size, crop_size=args.crop_size, multi_grid=args.multi_grid, multi_dilation=args.multi_dilation) #print(model) self.logger.info(model) # optimizer using different LR if not args.wo_backbone: params_list = [ { 'params': model.pretrained.parameters(), 'lr': args.lr }, ] else: params_list = [] 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) self.criterion = SegmentationMultiLosses(nclass=self.nclass) #self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux,nclass=self.nclass) self.model, self.optimizer = model, optimizer # using cuda if args.cuda: self.model = DataParallelModel(self.model).cuda() self.criterion = DataParallelCriterion(self.criterion).cuda() # finetune from a trained model if args.ft: args.start_epoch = 0 checkpoint = torch.load(args.ft_resume) if wo_head: print("WITHout HEAD !!!!!!!!!!") from collections import OrderedDict new = OrderedDict() for k, v in checkpoint['state_dict'].items(): if not k.startswith("head"): new[k] = v checkpoint['state_dict'] = new else: print("With HEAD !!!!!!!!!!") if args.cuda: self.model.module.load_state_dict(checkpoint['state_dict'], strict=False) else: self.model.load_state_dict(checkpoint['state_dict'], strict=False) # self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(args.ft_resume, checkpoint['epoch'])) # resuming checkpoint if args.resume: 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'] if args.cuda: self.model.module.load_state_dict(checkpoint['state_dict']) else: self.model.load_state_dict(checkpoint['state_dict']) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) self.best_pred = checkpoint['best_pred'] self.logger.info("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) # lr scheduler self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.trainloader), logger=self.logger, lr_step=args.lr_step) self.best_pred = 0.0
def test(args): # output folder outdir = args.save_folder if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # dataset testset = get_segmentation_dataset(args.dataset, split=args.split, mode=args.mode, transform=input_transform) # dataloader loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, dilated=args.dilated, multi_grid=args.multi_grid, stride=args.stride, lateral=args.lateral, jpu=args.jpu, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm, 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'])) # print(model) scales = [0.5, 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] if not args.ms: scales = [1.0] evaluator = MultiEvalModule(model, testset.num_class, scales=scales, flip=args.ms).cuda() evaluator.eval() metric = utils.SegmentationMetric(testset.num_class) tbar = tqdm(test_data) for i, (image, dst) in enumerate(tbar): if 'val' in args.mode: with torch.no_grad(): predicts = evaluator.parallel_forward(image) metric.update(dst, predicts) pixAcc, mIoU = metric.get() tbar.set_description('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU)) else: # with torch.no_grad(): # outputs = evaluator.parallel_forward(image) # predicts = [testset.make_pred(torch.max(output, 1)[1].cpu().numpy()) # for output in outputs] # for predict, impath in zip(predicts, dst): # mask = utils.get_mask_pallete(predict, args.dataset) # outname = os.path.splitext(impath)[0] + '.png' # mask.save(os.path.join(outdir, outname)) with torch.no_grad(): outputs = evaluator.parallel_forward(image) # predicts = [testset.make_pred(torch.max(output, 1)[1].cpu().numpy()) # for output in outputs] predicts = [ torch.softmax(output, 1).cpu().numpy() for output in outputs ] for predict, impath in zip(predicts, dst): # mask = utils.get_mask_pallete(predict, args.dataset) import numpy as np from PIL import Image mask = Image.fromarray( (predict[0, 1, :, :] * 255).astype(np.uint8)) outname = os.path.splitext(impath)[0] + '.bmp' mask.save(os.path.join(outdir, outname))
def test(args): # output folder outdir = '%s/msdanet_vis' % (args.dataset) if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), 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) else: # set split='test' for test set testset = get_segmentation_dataset(args.dataset, split='val', mode='vis', transform=input_transform) # dataloader loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm2d, base_size=args.base_size, crop_size=args.crop_size, multi_grid=args.multi_grid, multi_dilation=args.multi_dilation) # 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'], strict=False) print(model) num_class = testset.num_class evaluator = MultiEvalModule(model, testset.num_class, multi_scales=args.multi_scales).cuda() evaluator.eval() tbar = tqdm(test_data) def eval_batch(image, dst, evaluator, eval_mode): if eval_mode: # evaluation mode on validation set targets = dst outputs = evaluator.parallel_forward(image) batch_inter, batch_union, batch_correct, batch_label = 0, 0, 0, 0 for output, target in zip(outputs, targets): correct, labeled = utils.batch_pix_accuracy( output.data.cpu(), target) inter, union = utils.batch_intersection_union( output.data.cpu(), target, testset.num_class) batch_correct += correct batch_label += labeled batch_inter += inter batch_union += union return batch_correct, batch_label, batch_inter, batch_union else: # Visualize and dump the results im_paths = dst outputs = evaluator.parallel_forward(image) predicts = [ torch.max(output, 1)[1].cpu().numpy() + testset.pred_offset for output in outputs ] for predict, impath in zip(predicts, im_paths): mask = utils.get_mask_pallete(predict, args.dataset) outname = os.path.splitext(impath)[0] + '.png' mask.save(os.path.join(outdir, outname)) # dummy outputs for compatible with eval mode return 0, 0, 0, 0 total_inter, total_union, total_correct, total_label = \ np.int64(0), np.int64(0), np.int64(0), np.int64(0) for i, (image, dst) in enumerate(tbar): if torch_ver == "0.3": image = Variable(image, volatile=True) correct, labeled, inter, union = eval_batch( image, dst, evaluator, args.eval) else: with torch.no_grad(): correct, labeled, inter, union = eval_batch( image, dst, evaluator, args.eval) pixAcc, mIoU, IoU = 0, 0, 0 if args.eval: total_correct += correct.astype('int64') total_label += labeled.astype('int64') total_inter += inter.astype('int64') total_union += union.astype('int64') pixAcc = np.float64(1.0) * total_correct / ( np.spacing(1, dtype=np.float64) + total_label) IoU = np.float64(1.0) * total_inter / ( np.spacing(1, dtype=np.float64) + total_union) mIoU = IoU.mean() tbar.set_description('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU)) return pixAcc, mIoU, IoU, num_class
sys.path.append(r"E:\Project\PyTorch-Encoding\build\lib.win-amd64-3.6") import encoding from encoding.models import get_segmentation_model from encoding.nn import SyncBatchNorm from encoding.parallel import DataParallelModel from osgeo import gdal import threading gdal.AllRegister() # Get the model checkpoint = torch.load( r'E:\Project\PyTorch-Encoding\runs\arcs\deeplab\resnest269\model_best.pth.tar\model_best.pth.tar' ) model = get_segmentation_model("deeplab", dataset="arcs", backbone="resnest269", aux=True, se_loss=False, norm_layer=SyncBatchNorm, base_size=128, crop_size=128) model = DataParallelModel(model).cuda() model.module.load_state_dict(checkpoint['state_dict']) model.eval() def processData(tmpName): oriTileDir = "F:\\色林错\\dataSet\\" + str(tmpName) + r"\OriginTileData" # maskTileDir = "F:\\色林错\\dataSet\\" + str(tmpName) + r"\MaskTileData" tmpDir = "F:\\色林错\\dataSet\\" + str(tmpName) + r"\tmpTrainTest" if not os.path.exists(tmpDir): os.makedirs(tmpDir) length = dataSet[tmpName]["length"]
def test(args): # output folder outdir = 'outdir' if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # dataset data_kwargs = {'root': args.data_root} if args.eval: testset = get_segmentation_dataset(args.dataset, split='val', mode='testval', transform=input_transform, **data_kwargs) elif args.test_val: testset = get_segmentation_dataset(args.dataset, split='val', mode='test', transform=input_transform, **data_kwargs) else: testset = get_segmentation_dataset(args.dataset, split='test', mode='test', transform=input_transform, **data_kwargs) # dataloader loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) #model.base_size = args.base_size #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, 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'])) print(model) # 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] scales = [1.0] evaluator = MultiEvalModule(model, testset.num_class, scales=scales).cuda() evaluator.eval() metric = utils.SegmentationMetric(testset.num_class) tbar = tqdm(test_data) for i, (image, dst) in enumerate(tbar): if args.eval: with torch.no_grad(): predicts = evaluator.parallel_forward(image) metric.update(dst, predicts) pixAcc, mIoU = metric.get() tbar.set_description('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU)) else: with torch.no_grad(): outputs = evaluator.parallel_forward(image) predicts = [ testset.make_pred(torch.max(output, 1)[1].cpu().numpy()) for output in outputs ] for predict, impath in zip(predicts, dst): mask = utils.get_mask_pallete(predict, args.dataset) outname = os.path.splitext(impath)[0] + '.png' mask.save(os.path.join(outdir, outname))
def __init__(self, args): self.args = args # 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) testset = get_dataset(args.dataset, split='val', mode='val', **data_kwargs) self.train_sampler = torch.utils.data.distributed.DistributedSampler( trainset) self.val_sampler = torch.utils.data.distributed.DistributedSampler( testset) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} #self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, # drop_last=True, shuffle=True, **kwargs) #collate_fn=test_batchify_fn, self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size // args.world_size, drop_last=True, shuffle=False, sampler=self.train_sampler, **kwargs) #self.valloader = data.DataLoader(testset, batch_size=args.batch_size, self.valloader = data.DataLoader(testset, batch_size=args.test_batch_size // args.world_size, drop_last=False, shuffle=False, sampler=self.val_sampler, **kwargs) self.nclass = trainset.num_class #Norm_method = nn.SyncBatchNorm #Norm_method = nn.BatchNorm2d(momentum=0.01) Norm_method = nn.BatchNorm2d # model model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, multi_grid=args.multi_grid, se_loss=args.se_loss, norm_layer=Norm_method, lateral=args.lateral, root=args.backbone_path, base_size=args.base_size, crop_size=args.crop_size) if self.args.rank == 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) self.optimizer = optimizer #self.model = model # criterions self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux, nclass=self.nclass, se_weight=args.se_weight, aux_weight=args.aux_weight) device = torch.device('cuda:{}'.format(args.local_rank)) self.device = device # using cuda if args.cuda: #self.model = DataParallelModel(self.model).cuda() #self.model = self.model.cuda() sync_bn_model = FullModel(model, self.criterion) #self.model.cuda() #broadcast_params(self.model) #num_gpus = torch.cuda.device_count() #local_rank = args.local_rank % num_gpus #local_rank = args.local_rank #process_group = torch.distributed.new_group([args.local_rank]) #process_group = torch.distributed.new_group([args.rank]) #sync_bn_model = torch.nn.utils.convert_sync_batchnorm(self.model, process_group) sync_bn_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm( sync_bn_model) sync_bn_model = sync_bn_model.to(device) #self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[args.local_rank], output_device=args.local_rank) self.model = torch.nn.parallel.DistributedDataParallel( sync_bn_model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True) #self.criterion = DataParallelCriterion(self.criterion).cuda() #self.criterion = self.criterion.cuda() dist.barrier() # resuming checkpoint #if args.resume is not None and self.args.rank == 0: 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) old_state_dict = checkpoint['state_dict'] new_state_dict = dict() for k, v in old_state_dict.items(): if k.startswith('module.'): #new_state_dict[k[len('module.'):]] = old_state_dict[k] new_state_dict[k] = old_state_dict[k] else: new_state_dict[k] = old_state_dict[k] args.start_epoch = checkpoint['epoch'] if args.cuda: #self.model.module.load_state_dict(checkpoint['state_dict']) #self.model.load_state_dict(checkpoint['state_dict']) self.model.load_state_dict(new_state_dict) else: #self.model.load_state_dict(checkpoint['state_dict']) self.model.load_state_dict(new_state_dict) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) for state in self.optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.cuda() self.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 self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.trainloader), local_rank=self.args.rank) print('len(trainloader) : %.3f ' % (len(self.trainloader))) self.best_pred = 0.0 #for sumaryWriter self.track_loss = 0.0 self.track_pixAcc = 0.0 self.track_mIoU = 0.0
def __init__(self, args): self.args = args self.logger = utils.create_logger(self.args.exp_dir, "log") for k, v in vars(self.args).items(): self.logger.info((k, v)) if self.args.cuda: device = torch.device("cuda") self.logger.info("training on gpu:" + self.args.gpu_id) else: self.logger.info("training on cpu") device = torch.device("cpu") self.device = device #指定随机数 set_seed(args.random_seed) # args.log_name = str(args.checkname) #好像是可以加速 cudnn.benchmark = True #读取数据集,现在只有VOC self.logger.info('training on dataset ' + self.args.dataset) input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([.485, .456, .406], [.229, .224, .225]) ]) data_kwargs = { 'transform': input_transform, 'root': self.args.data_dir, 'base_size': self.args.base_size, 'crop_size': self.args.crop_size, 'logger': self.logger, 'scale': self.args.scale } trainset = get_segmentation_dataset(self.args.dataset, split='train', mode='train', **data_kwargs) testset = get_segmentation_dataset(self.args.dataset, split='val', mode='val', **data_kwargs) # dataloader kwargs = {'num_workers': self.args.num_workers, 'pin_memory': True} \ if self.args.cuda else {} self.train_iter = data.DataLoader(trainset, batch_size=self.args.batch_size, drop_last=True, shuffle=True, **kwargs) self.val_iter = data.DataLoader(testset, batch_size=self.args.batch_size, drop_last=False, shuffle=False, **kwargs) self.num_classes = trainset.num_classes self.input_channels = trainset.input_channels #create model kwargs = {'fuse_attention': self.args.fuse_attention} self.model = get_segmentation_model(args.arch, dataset=args.dataset, backbone=args.backbone) print("=> creating model %s" % self.args.arch) # self.model = archs.__dict__[self.args.arch](num_classes=self.num_classes, # input_channels=self.input_channels, **model_kwargs) self.model = self.model.to(device) # self.logger.info(self.model) self.optimizer = None params = filter(lambda p: p.requires_grad, self.model.parameters()) if self.args.optimizer == "Adam": self.optimizer = torch.optim.Adam( params, lr=self.args.lr, weight_decay=self.args.weight_decay) elif self.args.optimizer == 'SGD': self.optimizer = torch.optim.SGD( params, lr=self.args.lr, momentum=self.args.momentum, weight_decay=self.args.weight_decay) else: raise NotImplementedError #loss函数 self.criterion = nn.CrossEntropyLoss( ignore_index=trainset.IGNORE_INDEX) #语义分割评价指标 self.metric = SegmentationMetric(self.num_classes) #学习率策略 self.scheduler = LR_Scheduler(self.args.scheduler, base_lr = self.args.lr, num_epochs=self.args.epochs, \ iters_per_epoch=len(self.train_iter)) #创建实验结果保存目录 self.writer = SummaryWriter(args.exp_dir) # with open(os.path.join(args.exp_dir,'config.yml'), 'w') as f: # yaml.dump(config, f) #用tensoboard看一下模型结构 X, label = next(iter(self.train_iter)) self.writer.add_graph(self.model, X.to(device)) self.epoch_begin = 0 self.best_iou = 0.0 #在训练开始前看看输出是什么 val_log = self.validate(epoch=-1, is_visualize_segmentation=True) self.write_into_tensorboard(val_log, val_log, epoch=-1) #checkpoint_PATH if self.args.checkpoint_PATH is not None: if self.args.only_read_model: model, _, _, _, _ = load_checkpoint(model, self.args.checkpoint_PATH) else: model, self.epoch_begin, self.best_iou, self.optimizer = load_checkpoint( model, self.args.checkpoint_PATH, epoch_begin, best_iou, optimizer, scheduler)
def test(args): # output folder outdir = args.save_folder if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # dataset testset = get_segmentation_dataset(args.dataset, split=args.split, mode=args.mode, transform=input_transform) # dataloader loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, dilated=args.dilated, multi_grid=args.multi_grid, stride=args.stride, lateral=args.lateral, jpu=args.jpu, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm, 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'], strict=False) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) # print(model) scales = [0.5, 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] if not args.ms: scales = [1.0] evaluator = MultiEvalModule(model, testset.num_class, scales=scales, flip=args.ms).cuda() evaluator.eval() tbar = tqdm(test_data) total_inter, total_union, total_correct, total_label = 0, 0, 0, 0 result = [] for i, (image, dst) in enumerate(tbar): # print(dst) with torch.no_grad(): if i > 20: st = time.time() outputs = evaluator.forward(image[0].unsqueeze(0).cuda()) if i > 20: result.append(1 / (time.time() - st)) print(np.mean(result), np.std(result)) if 'val' in args.mode: # compute image IoU metric inter, union, area_pred, area_lab = batch_intersection_union( outputs, dst[0], testset.num_class) total_label += area_lab total_inter += inter total_union += union class_pixAcc = 1.0 * inter / (np.spacing(1) + area_lab) class_IoU = 1.0 * inter / (np.spacing(1) + union) print("img Classes pixAcc:", class_pixAcc) print("img Classes IoU:", class_IoU) else: # save prediction results predict = testset.make_pred( torch.max(output, 1)[1].cpu().numpy()) mask = utils.get_mask_pallete(predict, args.dataset) outname = os.path.splitext(dst[0])[0] + '.png' mask.save(os.path.join(outdir, outname)) if 'val' in args.mode: # compute set IoU metric pixAcc = 1.0 * total_inter / (np.spacing(1) + total_label) IoU = 1.0 * total_inter / (np.spacing(1) + total_union) mIoU = IoU.mean() print("set Classes pixAcc:", pixAcc) print("set Classes IoU:", IoU) print("set mean IoU:", mIoU)
def get_model(args): if args.backbone == "deeplabv2_multi": model = DeeplabMulti(num_classes=args.num_classes, pretrained=args.imagenet_pretrained) elif args.backbone == "new_deeplabv2_multi": model = NewDeeplabMulti(num_classes=args.num_classes, pretrained=args.imagenet_pretrained, use_se=args.use_se, train_bn=not args.freeze_bn, norm_style=args.norm_style) elif args.backbone == 'deeplabv3_resnest50': from encoding.models import get_segmentation_model from encoding.nn import SyncBatchNorm model = get_segmentation_model('deeplab', dataset='citys', backbone='resnest50', aux=True, norm_layer=SyncBatchNorm) elif args.backbone == 'deeplabv3_resnest101': from encoding.models import get_segmentation_model from encoding.nn import SyncBatchNorm model = get_segmentation_model('deeplab', dataset='citys', backbone='resnest101', aux=True, norm_layer=SyncBatchNorm) elif args.backbone == 'deeplabv2_vgg': model = DeeplabVGG(num_classes=args.num_classes) elif args.backbone == 'vgg16_fcn8s': model = VGG16_FCN8s(num_classes=args.num_classes) elif args.backbone == 'hrnet': raise NotImplementedError() # from graphs.models.hrnet import HighResolutionNet # model = HighResolutionNet(cfg) # model.init_weights(self.args.pretrained_ckpt_file) else: raise NotImplementedError() if 'deeplabv2' in args.backbone or 'vgg16_fcn8s' == args.backbone: params = model.optim_parameters(args) else: # https://github.com/zhanghang1989/PyTorch-Encoding/blob/master/experiments/segmentation/train.py#L153 params = [ { 'params': model.pretrained.parameters(), 'lr': args.lr }, ] if hasattr(model, 'head'): params.append({ 'params': model.head.parameters(), 'lr': args.lr * 10 }) print("Model head has 10x LR") if hasattr(model, 'auxlayer'): params.append({ 'params': model.auxlayer.parameters(), 'lr': args.lr * 10 }) print("Model auxlayer has 10x LR") args.numpy_transform = True return model, params
def test(args): # output folder outdir = 'outdir' if not os.path.exists(outdir): os.makedirs(outdir) # data transforms if not os.path.exists(args.save): os.makedirs(args.save) input_transform = transform.Compose([ transform.ToTensor(), 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) else: testset = get_segmentation_dataset(args.dataset, split='test', mode='test', transform=input_transform) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **kwargs) # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone = args.backbone, aux = args.aux, se_loss = args.se_loss, norm_layer = BatchNorm2d) # 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'])) print(model) evaluator = MultiEvalModuleCityscapes(model, testset.num_class).cuda() evaluator.eval() tbar = tqdm(test_data) interp = nn.Upsample(size=(1024, 2048), mode='bilinear') def eval_batch(image, dst, evaluator, eval_mode, hist, names): if eval_mode: # evaluation mode on validation set targets = dst outputs = evaluator.parallel_forward(image) batch_inter, batch_union, batch_correct, batch_label = 0, 0, 0, 0 for output, target, name in zip(outputs, targets, names): output = interp(output) correct, labeled = utils.batch_pix_accuracy(output.data.cpu(), target) inter, union = utils.batch_intersection_union( output.data.cpu(), target, testset.num_class) batch_correct += correct batch_label += labeled batch_inter += inter batch_union += union a = target.numpy().flatten() b = output.data.cpu() _, b = torch.max(b, 1) b = b.numpy().flatten() n = testset.num_class k = (a >= 0) & (a < n) hist += np.bincount(n * a[k].astype(int) + b[k], minlength = n ** 2).reshape(n, n) output = output.data.cpu().numpy()[0] output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) output_col = colorize_mask(output) output = Image.fromarray(output) name = name.split('/')[-1] output.save('%s/%s' % (args.save, name)) output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0])) return batch_correct, batch_label, batch_inter, batch_union, hist else: # test mode, dump the results im_paths = dst outputs = evaluator.parallel_forward(image) predicts = [torch.max(output, 1)[1].cpu().numpy() + testset.pred_offset for output in outputs] for predict, impath in zip(predicts, im_paths): mask = utils.get_mask_pallete(predict, args.dataset) outname = os.path.splitext(impath)[0] + '.png' mask.save(os.path.join(outdir, outname)) # dummy outputs for compatible with eval mode return 0, 0, 0, 0 total_inter, total_union, total_correct, total_label = \ np.int64(0), np.int64(0), np.int64(0), np.int64(0) hist =np.zeros((testset.num_class, testset.num_class)) for i, (image, dst, name) in enumerate(tbar): if torch_ver == "0.3": image = Variable(image, volatile=True) correct, labeled, inter, union, hist = eval_batch(image, dst, evaluator, args.eval, hist, name) else: with torch.no_grad(): correct, labeled, inter, union, hist = eval_batch(image, dst, evaluator, args.eval, hist, name) if args.eval: total_correct += correct.astype('int64') total_label += labeled.astype('int64') total_inter += inter.astype('int64') total_union += union.astype('int64') pixAcc = np.float64(1.0) * total_correct / (np.spacing(1, dtype=np.float64) + total_label) IoU = np.float64(1.0) * total_inter / (np.spacing(1, dtype=np.float64) + total_union) mIoU = IoU.mean() tbar.set_description( 'pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU)) mIoUs = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist)) print(str(round(np.nanmean(mIoUs) * 100, 2))) print(mIoUs)
def test(args): # output folder outdir = 'outdir' if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # dataset if args.eval: testset = get_segmentation_dataset(args.dataset, split='val', mode='val', transform=input_transform) else: testset = get_segmentation_dataset(args.dataset, split='test', mode='test', transform=input_transform) # dataloader loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm2d, 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 # pretrained_dict = checkpoint['state_dict'] # model_dict = model.state_dict() #for name, param in pretrained_dict.items(): # if name not in model_dict: # continue # if isinstance(param, Parameter): # backwards compatibility for serialized parameters # param = param.data # model_dict[name].copy_(param) #model.load_state_dict(checkpoint['state_dict']) # print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch'])) print(model) model = model.cuda() model.eval() run_time = list() for i in range(0, 100): input = torch.randn(1, 3, 512, 512).cuda() # ensure that context initialization and normal_() operations # finish before you start measuring time torch.cuda.synchronize() torch.cuda.synchronize() start = time.perf_counter() with torch.no_grad(): output = model(input) torch.cuda.synchronize() # wait for mm to finish end = time.perf_counter() print(end - start) run_time.append(end - start) run_time.pop(0) print('Mean running time is ', np.mean(run_time))
norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) print('begin') if __name__ == "__main__": print('please write the parse') args = Options().parse() print('i am here') pretrained = args.resume is None and args.verify is None if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=pretrained) model.base_size = args.base_size 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=torch.nn.BatchNorm2d if args.acc_bn else SyncBatchNorm, base_size=args.base_size, crop_size=args.crop_size) # model = encoding.models.get_segmentation_model(args.model, dataset=args.dataset, aux=args.aux, # backbone=args.backbone, # se_loss=args.se_loss, norm_layer=torch.nn.BatchNorm2d) print('Creating the model:') print(model) model.cuda() model.eval() # x = Variable(torch.Tensor(4, 3, 480, 480)).cuda() rgb = cv2.imread('example.jpg') rgb = np.transpose(rgb, (2, 0, 1))
def semseg(input_path, output_path=None, with_L0=False): """ param: input_path: str, path of input image output_path: str, path to save output image return: tuple, [animal_name, "background"] if pixels of "background" dominate, ["background", animal_name] else. """ sys.argv = sys.argv[:1] option = Options() args = option.parse() args.aux = True args.se_loss = True args.resume = "./checkpoints/encnet_jpu_res101_pcontext.pth.tar" # model checkpoint torch.manual_seed(args.seed) # data transforms input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # using L0_smooth to transform the orignal picture if with_L0: mid_result = os.path.join(os.path.dirname(input_path), "L0_result.png") L0_smooth(input_path, mid_result) input_path = mid_result # model model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, dilated=args.dilated, lateral=args.lateral, jpu=args.jpu, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm, 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, map_location=torch.device('cpu')) # strict=False, so that it is compatible with old pytorch saved models model.load_state_dict(checkpoint['state_dict'], strict=False) print("semseg model loaded successfully!") scales = [0.5, 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] if not args.ms: scales = [1.0] num_classes = datasets[args.dataset.lower()].NUM_CLASS evaluator = MultiEvalModule(model, num_classes, scales=scales, flip=args.ms).cuda() evaluator.eval() classes = np.array([ 'empty', 'aeroplane', 'bag', 'bed', 'bedclothes', 'bench', 'bicycle', 'bird', 'boat', 'book', 'bottle', 'building', 'bus', 'cabinet', 'car', 'cat', 'ceiling', 'chair', 'cloth', 'computer', 'cow', 'cup', 'curtain', 'dog', 'door', 'fence', 'floor', 'flower', 'food', 'grass', 'ground', 'horse', 'keyboard', 'light', 'motorbike', 'mountain', 'mouse', 'person', 'plate', 'platform', 'pottedplant', 'road', 'rock', 'sheep', 'shelves', 'sidewalk', 'sign', 'sky', 'snow', 'sofa', 'table', 'track', 'train', 'tree', 'truck', 'tvmonitor', 'wall', 'water', 'window', 'wood' ]) animals = ['bird', 'cat', 'cow', 'dog', 'horse', 'mouse', 'sheep'] img = input_transform(Image.open(input_path).convert('RGB')).unsqueeze(0) with torch.no_grad(): output = evaluator.parallel_forward(img)[0] predict = torch.max(output, 1)[1].cpu().numpy() + 1 pred_idx = np.unique(predict) pred_label = classes[pred_idx] print("[SemSeg] ", input_path, ": ", pred_label, sep='') main_pixels = 0 main_idx = -1 for idx, label in zip(pred_idx, pred_label): if label in animals: pixels = np.sum(predict == idx) if pixels > main_pixels: main_pixels = pixels main_idx = idx background_pixels = np.sum(predict != main_idx) main_animal = classes[main_idx] predict[predict != main_idx] = 29 mask_matrix = predict.copy() if output_path is not None: mask_matrix[np.where(mask_matrix != 29)] = 1 mask_matrix[np.where(mask_matrix == 29)] = 0 mask = utils.get_mask_pallete(mask_matrix, args.dataset) mask.save(output_path) if main_idx < 29: return predict, (main_animal, "background") else: return predict, ("background", main_animal)
def __init__(self, args): self.args = args # 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='ms_val' if args.multi_scale_eval else 'fast_val', **data_kwargs) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, **kwargs) if self.args.multi_scale_eval: kwargs['collate_fn'] = test_batchify_fn self.valloader = data.DataLoader(valset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, **kwargs) self.nclass = trainset.num_class # model if args.norm_layer == 'bn': norm_layer = BatchNorm2d elif args.norm_layer == 'sync_bn': assert args.multi_gpu, "SyncBatchNorm can only be used when multi GPUs are available!" norm_layer = SyncBatchNorm else: raise ValueError('Invalid norm_layer {}'.format(args.norm_layer)) model = get_segmentation_model( args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, se_loss=args.se_loss, norm_layer=norm_layer, base_size=args.base_size, crop_size=args.crop_size, multi_grid=True, multi_dilation=[2, 4, 8], only_pam=True, ) 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 }) if hasattr(model, 'auxlayer'): params_list.append({ 'params': model.auxlayer.parameters(), 'lr': args.lr }) optimizer = torch.optim.SGD(params_list, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # criterions self.criterion = SegmentationMultiLosses() self.model, self.optimizer = model, optimizer # using cuda if args.multi_gpu: self.model = DataParallelModel(self.model).cuda() self.criterion = DataParallelCriterion(self.criterion).cuda() else: self.model = self.model.cuda() self.criterion = self.criterion.cuda() self.single_device_model = self.model.module if self.args.multi_gpu else self.model # 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'] self.single_device_model.load_state_dict(checkpoint['state_dict']) if not args.ft and not (args.only_val or args.only_vis or args.only_infer): self.optimizer.load_state_dict(checkpoint['optimizer']) self.best_pred = checkpoint['best_pred'] print("=> loaded checkpoint '{}' (epoch {}), best_pred {}".format( args.resume, checkpoint['epoch'], checkpoint['best_pred'])) # clear start epoch if fine-tuning if args.ft: args.start_epoch = 0 # lr scheduler self.lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, 0.6) self.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(args.gpu, 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)
def __init__(self, args): self.args = args self.log_file = args.resume_dir + '/' + args.log_file input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # dataset data_kwargs = {'transform': input_transform, 'img_size': args.img_size} # trainset = get_segmentation_dataset(args.dataset, mode='train', augment=True, **data_kwargs) # valset = get_segmentation_dataset(args.dataset, mode='val', augment=False, **data_kwargs) testset = get_segmentation_dataset(args.testdata, mode='test', augment=False, whole_image=args.whole_image, **data_kwargs) test_vis_set = get_segmentation_dataset(args.testdata, mode='vis', augment=False, whole_image=args.whole_image, **data_kwargs) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} # self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, **kwargs) # self.valloader = data.DataLoader(valset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) self.testloader = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) self.visloader = data.DataLoader(test_vis_set, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) # self.nclass = trainset.NUM_CLASS # model model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, pretrained=args.pretrained, aux=args.aux, is_train=args.is_train, se_loss=args.se_loss, batchnorm=SynchronizedBatchNorm2d, img_size=args.img_size, dilated=args.dilated, deep_base=args.deep_base, multi_grid=args.multi_grid, multi_dilation=args.multi_dilation, ensemble=args.ensemble, resaux=args.resaux, aggaux=args.aggaux, output_stride=args.output_stride, test_size=args.test_img_size, high_rates=args.high_rates, aspp_rates=args.aspp_rates, aspp_out_dim=args.aspp_out_dim, agg_out_dim=args.agg_out_dim, up_conv=args.up_conv, refine=args.refine, refine_ver=args.refine_ver) # print(model) # optimizer using different LR # params_list = [{'params': model.pretrain_model.parameters(), 'lr': args.lr}, # {'params': model.head.parameters(), 'lr': args.lr * args.head_lr_factor}] # optimizer = torch.optim.SGD(params_list, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # optimizer = torch.optim.Adam(params_list, lr=args.lr, betas=(0.9, 0.999), weight_decay=args.weight_decay) # clear start epoch if fine-turning # if args.ft: # args.start_epoch = 0 # criterions ''' self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux, nclass=self.nclass, xentropy_weight=args.xentropy_weight, mae_weight=args.mae_weight, dice_weight=args.dice_weight, lovasz_weight=args.lovasz_weight, loss1_4=args.loss1_4, lov1_4=args.lov1_4, margin_weight=args.margin_weight, margin=args.margin, k=args.k, c=args.c) self.model, self.optimizer = model, optimizer ''' self.model = model # using cuda # cudnn.benchmark = True self.model = torch.nn.DataParallel(self.model).cuda() # self.criterion = self.criterion.cuda() self.sigmoid = nn.Sigmoid().cuda() # lr_scheduler # self.scheduler = utils.LR_Scheduler(args.lr_mode, args.lr, args.epochs, len(self.trainloader)) self.best_pred = 0.0 # resuming chechpoint if args.resume_dir is not None: if not os.path.isfile(args.resume_dir + '/checkpoint.pth.tar'): print('=> no chechpoint found at {}'.format(args.resume_dir)) logger(self.log_file, '=> no chechpoint found at {}'.format(args.resume_dir)) else: checkpoint = torch.load(args.resume_dir + '/checkpoint.pth.tar') args.start_epoch = checkpoint['epoch'] self.model.module.load_state_dict(checkpoint['state_dict']) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) self.best_pred = checkpoint['best_pred'] print('=> loaded checkpoint {0} (epoch {1})'.format( args.resume_dir, checkpoint['epoch'])) logger( self.log_file, '=> loaded checkpoint {0} (epoch {1})'.format( args.resume_dir, checkpoint['epoch']))
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) # optimizer = torch.optim.Adam(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']) ''' checkpoint = torch.load(args.resume, map_location='cpu') args.start_epoch = checkpoint['epoch'] model.module.load_state_dict(checkpoint['state_dict']) model.cuda() ''' 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)) # train_losses = [2.855, 2.513, 2.275, 2.128, 2.001, 1.875, 1.855, 1.916, 1.987, 1.915, 1.952] train_losses = [] def training(epoch): train_sampler.set_epoch(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))) train_losses.append(train_loss / len(trainloader)) if epoch > 1: if train_losses[epoch] < train_losses[epoch - 1]: utils.save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'best_pred': new_preds[(epoch - 1) // 10], }, args, False, filename='checkpoint_train.pth.tar') plt.plot(train_losses) plt.xlabel('Epoch') plt.ylabel('Train_loss') plt.title('Train_Loss') plt.grid() plt.savefig('./loss_fig/train_losses.pdf') plt.savefig('./loss_fig/train_losses.svg') plt.close() # p_m = [(0.3, 0.05), (0.23, 0.54)] # new_preds = [0.175, 0.392] p_m = [] new_preds = [] 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(): pred = model(image)[0] target = target.cuda(args.gpu) metric.update(target, pred) if i % 100 == 0: all_metircs = metric.get_all() all_metircs = utils.torch_dist_sum(args.gpu, *all_metircs) pixAcc, mIoU = utils.get_pixacc_miou(*all_metircs) if args.gpu == 0: print('pixAcc: %.3f, mIoU1: %.3f' % (pixAcc, mIoU)) all_metircs = metric.get_all() all_metircs = utils.torch_dist_sum(args.gpu, *all_metircs) pixAcc, mIoU = utils.get_pixacc_miou(*all_metircs) if args.gpu == 0: print('pixAcc: %.3f, mIoU2: %.3f' % (pixAcc, mIoU)) p_m.append((pixAcc, mIoU)) plt.plot(p_m) plt.xlabel('10 Epoch') plt.ylabel('pixAcc, mIoU') plt.title('pixAcc, mIoU') plt.grid() plt.legend(('pixAcc', 'mIoU')) plt.savefig('./loss_fig/pixAcc_mIoU.pdf') plt.savefig('./loss_fig/pixAcc_mIoU.svg') plt.close() if args.eval: return new_pred = (pixAcc + mIoU) / 2 new_preds.append(new_pred) plt.plot(new_preds) plt.xlabel('10 Epoch') plt.ylabel('new_predication') plt.title('new_predication') plt.grid() plt.savefig('./loss_fig/new_predication.pdf') plt.savefig('./loss_fig/new_predication.svg') plt.close() 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, filename='checkpoint_train_{}.pth.tar'.format(epoch + 1)) if args.export: if args.gpu == 0: torch.save(model.module.state_dict(), args.export + '.pth') return if args.eval: validation(args.start_epoch) return 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 or epoch == args.epochs - 1: validation(epoch) elapsed = time.time() - tic if args.gpu == 0: print(f'Epoch: {epoch}, Time cost: {elapsed}') validation(epoch)
def __init__(self, args): self.args = args # 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_segmentation_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs) testset = get_segmentation_dataset(args.dataset, split='val', mode='val', **data_kwargs) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=False, shuffle=True, **kwargs) self.valloader = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) self.nclass = trainset.num_class # model model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, dilated=args.dilated, multi_grid=args.multi_grid, stride=args.stride, lateral=args.lateral, jpu=args.jpu, aux=args.aux, se_loss=args.se_loss, norm_layer=SyncBatchNorm, base_size=args.base_size, crop_size=args.crop_size) # print(model) # optimizer using different LR params_list = [ { 'params': model.pretrained.parameters(), 'lr': args.lr }, ] if hasattr(model, 'jpu') and model.jpu: params_list.append({ 'params': model.jpu.parameters(), 'lr': args.lr * 10 }) 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) class_balance_weight = 'None' if args.dataset == "pcontext60": class_balance_weight = torch.tensor([ 1.3225e-01, 2.0757e+00, 1.8146e+01, 5.5052e+00, 2.2060e+00, 2.8054e+01, 2.0566e+00, 1.8598e+00, 2.4027e+00, 9.3435e+00, 3.5990e+00, 2.7487e-01, 1.4216e+00, 2.4986e+00, 7.7258e-01, 4.9020e-01, 2.9067e+00, 1.2197e+00, 2.2744e+00, 2.0444e+01, 3.0057e+00, 1.8167e+01, 3.7405e+00, 5.6749e-01, 3.2631e+00, 1.5007e+00, 5.5519e-01, 1.0056e+01, 1.8952e+01, 2.6792e-01, 2.7479e-01, 1.8309e+00, 2.0428e+01, 1.4788e+01, 1.4908e+00, 1.9113e+00, 2.6166e+02, 2.3233e-01, 1.9096e+01, 6.7025e+00, 2.8756e+00, 6.8804e-01, 4.4140e+00, 2.5621e+00, 4.4409e+00, 4.3821e+00, 1.3774e+01, 1.9803e-01, 3.6944e+00, 1.0397e+00, 2.0601e+00, 5.5811e+00, 1.3242e+00, 3.0088e-01, 1.7344e+01, 2.1569e+00, 2.7216e-01, 5.8731e-01, 1.9956e+00, 4.4004e+00 ]) elif args.dataset == "ade20k": class_balance_weight = torch.tensor([ 0.0772, 0.0431, 0.0631, 0.0766, 0.1095, 0.1399, 0.1502, 0.1702, 0.2958, 0.3400, 0.3738, 0.3749, 0.4059, 0.4266, 0.4524, 0.5725, 0.6145, 0.6240, 0.6709, 0.6517, 0.6591, 0.6818, 0.9203, 0.9965, 1.0272, 1.0967, 1.1202, 1.2354, 1.2900, 1.5038, 1.5160, 1.5172, 1.5036, 2.0746, 2.1426, 2.3159, 2.2792, 2.6468, 2.8038, 2.8777, 2.9525, 2.9051, 3.1050, 3.1785, 3.3533, 3.5300, 3.6120, 3.7006, 3.6790, 3.8057, 3.7604, 3.8043, 3.6610, 3.8268, 4.0644, 4.2698, 4.0163, 4.0272, 4.1626, 4.3702, 4.3144, 4.3612, 4.4389, 4.5612, 5.1537, 4.7653, 4.8421, 4.6813, 5.1037, 5.0729, 5.2657, 5.6153, 5.8240, 5.5360, 5.6373, 6.6972, 6.4561, 6.9555, 7.9239, 7.3265, 7.7501, 7.7900, 8.0528, 8.5415, 8.1316, 8.6557, 9.0550, 9.0081, 9.3262, 9.1391, 9.7237, 9.3775, 9.4592, 9.7883, 10.6705, 10.2113, 10.5845, 10.9667, 10.8754, 10.8274, 11.6427, 11.0687, 10.8417, 11.0287, 12.2030, 12.8830, 12.5082, 13.0703, 13.8410, 12.3264, 12.9048, 12.9664, 12.3523, 13.9830, 13.8105, 14.0345, 15.0054, 13.9801, 14.1048, 13.9025, 13.6179, 17.0577, 15.8351, 17.7102, 17.3153, 19.4640, 17.7629, 19.9093, 16.9529, 19.3016, 17.6671, 19.4525, 20.0794, 18.3574, 19.1219, 19.5089, 19.2417, 20.2534, 20.0332, 21.7496, 21.5427, 20.3008, 21.1942, 22.7051, 23.3359, 22.4300, 20.9934, 26.9073, 31.7362, 30.0784 ]) elif args.dataset == "cocostuff": class_balance_weight = torch.tensor([ 4.8557e-02, 6.4709e-02, 3.9255e+00, 9.4797e-01, 1.2703e+00, 1.4151e+00, 7.9733e-01, 8.4903e-01, 1.0751e+00, 2.4001e+00, 8.9736e+00, 5.3036e+00, 6.0410e+00, 9.3285e+00, 1.5952e+00, 3.6090e+00, 9.8772e-01, 1.2319e+00, 1.9194e+00, 2.7624e+00, 2.0548e+00, 1.2058e+00, 3.6424e+00, 2.0789e+00, 1.7851e+00, 6.7138e+00, 2.1315e+00, 6.9813e+00, 1.2679e+02, 2.0357e+00, 2.2933e+01, 2.3198e+01, 1.7439e+01, 4.1294e+01, 7.8678e+00, 4.3444e+01, 6.7543e+01, 1.0066e+01, 6.7520e+00, 1.3174e+01, 3.3499e+00, 6.9737e+00, 2.1482e+00, 1.9428e+01, 1.3240e+01, 1.9218e+01, 7.6836e-01, 2.6041e+00, 6.1822e+00, 1.4070e+00, 4.4074e+00, 5.7792e+00, 1.0321e+01, 4.9922e+00, 6.7408e-01, 3.1554e+00, 1.5832e+00, 8.9685e-01, 1.1686e+00, 2.6487e+00, 6.5354e-01, 2.3801e-01, 1.9536e+00, 1.5862e+00, 1.7797e+00, 2.7385e+01, 1.2419e+01, 3.9287e+00, 7.8897e+00, 7.5737e+00, 1.9758e+00, 8.1962e+01, 3.6922e+00, 2.0039e+00, 2.7333e+00, 5.4717e+00, 3.9048e+00, 1.9184e+01, 2.2689e+00, 2.6091e+02, 4.7366e+01, 2.3844e+00, 8.3310e+00, 1.4857e+01, 6.5076e+00, 2.0854e-01, 1.0425e+00, 1.7386e+00, 1.1973e+01, 5.2862e+00, 1.7341e+00, 8.6124e-01, 9.3702e+00, 2.8545e+00, 6.0123e+00, 1.7560e-01, 1.8128e+00, 1.3784e+00, 1.3699e+00, 2.3728e+00, 6.2819e-01, 1.3097e+00, 4.7892e-01, 1.0268e+01, 1.2307e+00, 5.5662e+00, 1.2867e+00, 1.2745e+00, 4.7505e+00, 8.4029e+00, 1.8679e+00, 1.0519e+01, 1.1240e+00, 1.4975e-01, 2.3146e+00, 4.1265e-01, 2.5896e+00, 1.4537e+00, 4.5575e+00, 7.8143e+00, 1.4603e+01, 2.8812e+00, 1.8868e+00, 7.8131e+01, 1.9323e+00, 7.4980e+00, 1.2446e+01, 2.1856e+00, 3.0973e+00, 4.1270e-01, 4.9016e+01, 7.1001e-01, 7.4035e+00, 2.3395e+00, 2.9207e-01, 2.4156e+00, 3.3211e+00, 2.1300e+00, 2.4533e-01, 1.7081e+00, 4.6621e+00, 2.9199e+00, 1.0407e+01, 7.6207e-01, 2.7806e-01, 3.7711e+00, 1.1852e-01, 8.8280e+00, 3.1700e-01, 6.3765e+01, 6.6032e+00, 5.2177e+00, 4.3596e+00, 6.2965e-01, 1.0207e+00, 1.1731e+01, 2.3935e+00, 9.2767e+00, 1.1023e-01, 3.6947e+00, 1.3943e+00, 2.3407e+00, 1.2112e-01, 2.8518e+00, 2.8195e+00, 1.0078e+00, 1.6614e+00, 6.5307e-01, 1.9070e+01, 2.7231e+00, 6.0769e-01 ]) # criterions self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux, nclass=self.nclass, se_weight=args.se_weight, aux_weight=args.aux_weight, weight=class_balance_weight) self.model, self.optimizer = model, optimizer # using cuda if args.cuda: self.model = DataParallelModel(self.model).cuda() self.criterion = DataParallelCriterion(self.criterion).cuda() # resuming checkpoint self.best_pred = 0.0 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'] if args.cuda: self.model.module.load_state_dict(checkpoint['state_dict']) else: self.model.load_state_dict(checkpoint['state_dict']) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) self.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 self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.trainloader))
from encoding.models import get_model, get_segmentation_model, MultiEvalModule from lcd import lcdutils from option import Options args = Options().parse() # Get the model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True).cuda() else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm2d, base_size=args.base_size, crop_size=args.crop_size, multi_grid=args.multi_grid, multi_dilation=args.multi_dilation).cuda() # resuming checkpoint if args.resume_dir is not None: args.resume = os.path.join(args.resume_dir, "DANet101.pth.tar") 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'], strict=False) model.eval()
def __init__(self, args): self.args = args args.log_name = str(args.checkname) self.logger = utils.create_logger(args.log_root, args.log_name) # data transforms input_transform = None # dataset data_kwargs = { 'transform': input_transform, 'base_size': args.base_size, 'crop_size': args.crop_size } trainset = get_dataset(args.dataset, split='train', 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 {} self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, **kwargs) self.valloader = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) self.nclass = trainset.num_class # model model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm2d, base_size=args.base_size, crop_size=args.crop_size) #print(model) self.logger.info(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) self.criterion = SegmentationMultiLosses(nclass=self.nclass) #self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux,nclass=self.nclass) self.model, self.optimizer = model, optimizer # using cuda if args.cuda: self.model = torch.nn.DataParallel(self.model).cuda() self.criterion = torch.nn.DataParallel(self.criterion).cuda() # finetune from a trained model if args.ft: args.start_epoch = 0 checkpoint = torch.load(args.ft_resume) if args.cuda: self.model.module.load_state_dict(checkpoint, strict=False) else: self.model.load_state_dict(checkpoint, strict=False) self.logger.info("=> loaded checkpoint '{}' (epoch {})".format( args.ft_resume, args.start_epoch)) # resuming checkpoint if args.resume: 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'] if args.cuda: self.model.module.load_state_dict(checkpoint['state_dict']) else: self.model.load_state_dict(checkpoint['state_dict']) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) self.best_pred = checkpoint['best_pred'] self.logger.info("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) # lr scheduler self.scheduler = utils.LR_Scheduler( args.lr_scheduler, args.lr, args.epochs, len(self.trainloader), # logger=self.logger, lr_step=args.lr_step) self.best_pred = 0.0
def test(args): # output folder outdir = 'outdir' if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # dataset if args.eval: testset = get_segmentation_dataset(args.dataset, split='val', mode='val', transform=input_transform, return_file=True) else: testset = get_segmentation_dataset(args.dataset, split='test', mode='test', transform=input_transform) # dataloader loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm2d, 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 pretrained_dict = checkpoint['state_dict'] model_dict = model.state_dict() for name, param in pretrained_dict.items(): if name not in model_dict: continue if isinstance(param, Parameter): # backwards compatibility for serialized parameters param = param.data model_dict[name].copy_(param) #model.load_state_dict(checkpoint['state_dict']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) print(model) # count parameter number pytorch_total_params = sum(p.numel() for p in model.parameters()) print("Total number of parameters: %d" % pytorch_total_params) evaluator = MultiEvalModule(model, testset.num_class).cuda() evaluator.eval() tbar = tqdm(test_data) def eval_batch(image, dst, im_paths, evaluator, eval_mode): if eval_mode: # evaluation mode on validation set targets = dst outputs = evaluator.parallel_forward(image) batch_inter, batch_union, batch_correct, batch_label = 0, 0, 0, 0 for output, target in zip(outputs, targets): correct, labeled = utils.batch_pix_accuracy( output.data.cpu(), target) inter, union = utils.batch_intersection_union( output.data.cpu(), target, testset.num_class) batch_correct += correct batch_label += labeled batch_inter += inter batch_union += union # save outputs predicts = [ torch.max(output, 1)[1].cpu().numpy() # + testset.pred_offset for output in outputs ] for predict, impath, target in zip(predicts, im_paths, targets): mask = utils.get_mask_pallete(predict, args.dataset) outname = os.path.splitext(impath)[0] + '.png' mask.save(os.path.join(outdir, outname)) # save ground truth into png format target = target.data.cpu().numpy() target = utils.get_mask_pallete(target, args.dataset) outname = os.path.splitext(impath)[0] + '_gtruth.png' target.save(os.path.join(outdir, outname)) return batch_correct, batch_label, batch_inter, batch_union else: # test mode, dump the results im_paths = dst outputs = evaluator.parallel_forward(image) predicts = [ torch.max(output, 1)[1].cpu().numpy() # + testset.pred_offset for output in outputs ] for predict, impath in zip(predicts, im_paths): mask = utils.get_mask_pallete(predict, args.dataset) outname = os.path.splitext(impath)[0] + '.png' mask.save(os.path.join(outdir, outname)) # dummy outputs for compatible with eval mode return 0, 0, 0, 0 total_inter, total_union, total_correct, total_label = \ np.int64(0), np.int64(0), np.int64(0), np.int64(0) for i, (image, dst, img_paths) in enumerate(tbar): if torch_ver == "0.3": image = Variable(image, volatile=True) correct, labeled, inter, union = eval_batch( image, dst, img_paths, evaluator, args.eval) else: with torch.no_grad(): correct, labeled, inter, union = eval_batch( image, dst, img_paths, evaluator, args.eval) if args.eval: total_correct += correct.astype('int64') total_label += labeled.astype('int64') total_inter += inter.astype('int64') total_union += union.astype('int64') pixAcc = np.float64(1.0) * total_correct / ( np.spacing(1, dtype=np.float64) + total_label) IoU = np.float64(1.0) * total_inter / ( np.spacing(1, dtype=np.float64) + total_union) mIoU = IoU.mean() tbar.set_description('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU))
def test(args): directory = "runs/val_summary/%s/%s/%s/" % (args.dataset, args.model, args.resume) if not os.path.exists(directory): os.makedirs(directory) writer = SummaryWriter(directory) # output folder outdir = 'outdir' if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), 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) elif args.test_val: testset = get_segmentation_dataset(args.dataset, split='val', mode='test', transform=input_transform) else: testset = get_segmentation_dataset(args.dataset, split='test', mode='test', transform=input_transform) # dataloader loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) Norm_method = torch.nn.BatchNorm2d # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) #model.base_size = args.base_size #model.crop_size = args.crop_size else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, multi_grid=args.multi_grid, num_center=args.num_center, norm_layer=Norm_method, root=args.backbone_path, 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.module.load_state_dict(checkpoint['state_dict']) old_state_dict = checkpoint['state_dict'] new_state_dict = dict() for k, v in old_state_dict.items(): if k.startswith('module.'): #new_state_dict[k[len('module.'):]] = old_state_dict[k] new_state_dict[k[len('model.module.'):]] = old_state_dict[k] #new_state_dict[k] = old_state_dict[k] else: new_state_dict[k] = old_state_dict[k] #new_k = 'module.' + k #new_state_dict[new_k] = old_state_dict[k] model.load_state_dict(new_state_dict) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) print(model) scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \ [0.75, 1.0, 1.25, 1.5, 1.75, 2.0] if args.dataset == 'ade20k': scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0] if not args.ms: scales = [1.0] if args.dataset == 'ade20k': evaluator = MultiEvalModule2(model, testset.num_class, scales=scales, flip=args.ms).cuda() else: evaluator = MultiEvalModule(model, testset.num_class, scales=scales, flip=args.ms).cuda() evaluator.eval() metric = utils.SegmentationMetric(testset.num_class) tbar = tqdm(test_data) for i, (image, dst) in enumerate(tbar): if args.eval: with torch.no_grad(): predicts = evaluator.parallel_forward(image) metric.update(dst, predicts) pixAcc, mIoU = metric.get() tbar.set_description('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU)) writer.add_scalar('pixAcc', pixAcc, i) writer.add_scalar('mIoU', mIoU, i) else: with torch.no_grad(): outputs = evaluator.parallel_forward(image) predicts = [ testset.make_pred(torch.max(output, 1)[1].cpu().numpy()) for output in outputs ] for predict, impath in zip(predicts, dst): mask = utils.get_mask_pallete(predict, args.dataset) outname = os.path.splitext(impath)[0] + '.png' mask.save(os.path.join(outdir, outname)) writer.close()
def __init__(self, args): if args.se_loss: args.checkname = args.checkname + "_se" self.args = args # 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_segmentation_dataset(args.dataset, split='train', mode='train', **data_kwargs) testset = get_segmentation_dataset(args.dataset, split='val', mode ='val', **data_kwargs) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': False} \ if args.cuda else {} self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, **kwargs) self.valloader = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) self.nclass = trainset.num_class # model model = get_segmentation_model(args.model, dataset=args.dataset, backbone = args.backbone, aux = args.aux, se_loss = args.se_loss, norm_layer = BatchNorm2d, base_size=args.base_size, crop_size=args.crop_size) print(model) # count parameter number pytorch_total_params = sum(p.numel() for p in model.parameters()) print("Total number of parameters: %d"%pytorch_total_params) # optimizer using different LR params_list = [{'params': model.pretrained.parameters(), 'lr': args.lr},] if hasattr(model, 'head'): if args.diflr: params_list.append({'params': model.head.parameters(), 'lr': args.lr*10}) else: params_list.append({'params': model.head.parameters(), 'lr': args.lr}) if hasattr(model, 'auxlayer'): if args.diflr: params_list.append({'params': model.auxlayer.parameters(), 'lr': args.lr*10}) else: params_list.append({'params': model.auxlayer.parameters(), 'lr': args.lr}) optimizer = torch.optim.SGD(params_list, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) #optimizer = torch.optim.ASGD(params_list, # lr=args.lr, # weight_decay=args.weight_decay) # criterions self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux, nclass=self.nclass) self.model, self.optimizer = model, optimizer # using cuda if args.cuda: self.model = DataParallelModel(self.model).cuda() self.criterion = DataParallelCriterion(self.criterion).cuda() # resuming checkpoint if args.resume is not None and len(args.resume)>0: 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'] if args.cuda: # load weights for the same model # self.model.module.load_state_dict(checkpoint['state_dict']) # model and checkpoint have different strucutures pretrained_dict = checkpoint['state_dict'] model_dict = self.model.module.state_dict() for name, param in pretrained_dict.items(): if name not in model_dict: continue if isinstance(param, Parameter): # backwards compatibility for serialized parameters param = param.data model_dict[name].copy_(param) else: self.model.load_state_dict(checkpoint['state_dict']) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) self.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 self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.trainloader),lr_step=args.lr_step) self.best_pred = 0.0
def __init__(self, args): self.args = args args.log_name = str(args.checkname) self.logger = utils.create_logger(args.log_root, args.log_name) # 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, 'logger': self.logger, 'scale': args.scale } trainset = get_segmentation_dataset(args.dataset, split='train', mode='train', **data_kwargs) testset = get_segmentation_dataset(args.dataset, split='val', mode='val', **data_kwargs) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, **kwargs) self.valloader = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) self.nclass = trainset.num_class # model model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm2d, base_size=args.base_size, crop_size=args.crop_size, multi_grid=args.multi_grid, multi_dilation=args.multi_dilation) #print(model) self.logger.info(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 }) cityscape_weight = torch.FloatTensor([ 0.8373, 0.918, 0.866, 1.0345, 1.0166, 0.9969, 0.9754, 1.0489, 0.8786, 1.0023, 0.9539, 0.9843, 1.1116, 0.9037, 1.0865, 1.0955, 1.0865, 1.1529, 1.0507 ]) optimizer = torch.optim.SGD(params_list, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) #weight for class imbalance # self.criterion = SegmentationMultiLosses(nclass=self.nclass, weight=cityscape_weight) self.criterion = SegmentationMultiLosses(nclass=self.nclass) #self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux,nclass=self.nclass) self.model, self.optimizer = model, optimizer # using cuda if args.cuda: self.model = DataParallelModel(self.model).cuda() self.criterion = DataParallelCriterion(self.criterion).cuda() # finetune from a trained model if args.ft: args.start_epoch = 0 checkpoint = torch.load(args.ft_resume) if args.cuda: self.model.module.load_state_dict(checkpoint['state_dict'], strict=False) else: self.model.load_state_dict(checkpoint['state_dict'], strict=False) self.logger.info("=> loaded checkpoint '{}' (epoch {})".format( args.ft_resume, checkpoint['epoch'])) # resuming checkpoint if args.resume: 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'] if args.cuda: self.model.module.load_state_dict(checkpoint['state_dict']) else: self.model.load_state_dict(checkpoint['state_dict']) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) self.best_pred = checkpoint['best_pred'] self.logger.info("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) # lr scheduler self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.trainloader), logger=self.logger, lr_step=args.lr_step) self.best_pred = 0.0
def test(args): # output folder outdir = 'outdir' if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # dataset if args.eval: testset = get_dataset(args.dataset, split='val', mode='testval', transform=input_transform) elif args.test_val: testset = get_dataset(args.dataset, split='val', mode='test', transform=input_transform) else: 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 {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) # model pretrained = args.resume is None and args.verify is None if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=pretrained) model.base_size = args.base_size model.crop_size = args.crop_size else: # my model_kwargs = {} if args.choice_indices is not None: assert 'alone_resnest50' in args.backbone model_kwargs['choice_indices'] = args.choice_indices # model = get_segmentation_model( args.model, dataset=args.dataset, 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, **model_kwargs) # resuming checkpoint 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, map_location='cpu')) elif args.resume is not None and os.path.isfile(args.resume): checkpoint = torch.load(args.resume, map_location='cpu') # strict=False, so that it is compatible with old pytorch saved models model.load_state_dict(checkpoint['state_dict']) print("=> loaded checkpoint '{}'".format(args.resume)) elif not pretrained: raise RuntimeError("=> no checkpoint found") print(model) if args.acc_bn: from encoding.utils.precise_bn import update_bn_stats 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(ReturnFirstClosure(trainset), batch_size=args.batch_size, drop_last=True, shuffle=True, **loader_kwargs) print('Reseting BN statistics') #model.apply(reset_bn_statistics) model.cuda() update_bn_stats(model, trainloader) 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 evaluator = MultiEvalModule(model, testset.num_class, scales=scales).cuda() evaluator.eval() metric = utils.SegmentationMetric(testset.num_class) tbar = tqdm(test_data) for i, (image, dst) in enumerate(tbar): if args.eval: with torch.no_grad(): predicts = evaluator.parallel_forward(image) metric.update(dst, predicts) pixAcc, mIoU = metric.get() tbar.set_description('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU)) else: with torch.no_grad(): outputs = evaluator.parallel_forward(image) predicts = [ testset.make_pred(torch.max(output, 1)[1].cpu().numpy()) for output in outputs ] for predict, impath in zip(predicts, dst): mask = utils.get_mask_pallete(predict, args.dataset) outname = os.path.splitext(impath)[0] + '.png' mask.save(os.path.join(outdir, outname)) if args.eval: print('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU))
def __init__(self, args): self.args = args args.log_name = str(args.checkname) args.log_root = os.path.join(args.dataset, args.log_root) # dataset/log/ self.logger = utils.create_logger(args.log_root, args.log_name) # 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, 'logger': self.logger, 'scale': args.scale} trainset = get_segmentation_dataset(args.dataset, split='trainval', mode='trainval', **data_kwargs) testset = get_segmentation_dataset(args.dataset, split='val', mode='val', # crop fixed size as model input **data_kwargs) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, **kwargs) self.valloader = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) self.nclass = trainset.num_class # model model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, norm_layer=BatchNorm2d, base_size=args.base_size, crop_size=args.crop_size, ) #print(model) self.logger.info(model) # optimizer using different LR params_list = [{'params': model.pretrained.parameters(), 'lr': args.lr},] if hasattr(model, 'head'): print("this model has object, head") params_list.append({'params': model.head.parameters(), 'lr': args.lr*10}) optimizer = torch.optim.SGD(params_list, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) self.criterion = SegmentationLosses(nclass=self.nclass) self.model, self.optimizer = model, optimizer # using cuda if args.cuda: self.model = DataParallelModel(self.model).cuda() self.criterion = DataParallelCriterion(self.criterion).cuda() # resuming checkpoint if args.resume: 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'] if args.cuda: self.model.module.load_state_dict(checkpoint['state_dict']) else: self.model.load_state_dict(checkpoint['state_dict']) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) self.best_pred = checkpoint['best_pred'] self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch'])) # lr scheduler self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.trainloader), logger=self.logger, lr_step=args.lr_step) self.best_pred = 0.0
def __init__(self, args): self.args = args # 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_segmentation_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs) testset = get_segmentation_dataset(args.dataset, split='val', mode='val', **data_kwargs) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, **kwargs) self.valloader = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs) self.nclass = trainset.num_class # model model = get_segmentation_model( args.model, dataset=args.dataset, backbone=args.backbone, dilated=args.dilated, lateral=args.lateral, jpu=args.jpu, aux=args.aux, se_loss=args.se_loss, norm_layer=torch.nn.BatchNorm2d, ## BatchNorm2d base_size=args.base_size, crop_size=args.crop_size) print(model) # optimizer using different LR params_list = [ { 'params': model.pretrained.parameters(), 'lr': args.lr }, ] if hasattr(model, 'jpu'): params_list.append({ 'params': model.jpu.parameters(), 'lr': args.lr * 10 }) 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 self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux, nclass=self.nclass, se_weight=args.se_weight, aux_weight=args.aux_weight) self.model, self.optimizer = model, optimizer # using cuda if args.cuda: self.model = DataParallelModel(self.model).cuda() self.criterion = DataParallelCriterion(self.criterion).cuda() # 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'] if args.cuda: self.model.module.load_state_dict(checkpoint['state_dict']) else: self.model.load_state_dict(checkpoint['state_dict']) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) self.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 self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.trainloader)) self.best_pred = 0.0
def test(args): # output folder outdir = args.save_folder if not os.path.exists(outdir): os.makedirs(outdir) # data transforms input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) # dataset testset = get_segmentation_dataset(args.dataset, split=args.split, mode=args.mode, transform=input_transform) # dataloader loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \ if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.test_batch_size, drop_last=False, shuffle=False, collate_fn=test_batchify_fn, **loader_kwargs) # model if args.model_zoo is not None: model = get_model(args.model_zoo, pretrained=True) else: model = get_segmentation_model(args.model, dataset=args.dataset, backbone=args.backbone, dilated=args.dilated, multi_grid=args.multi_grid, stride=args.stride, lateral=args.lateral, jpu=args.jpu, aux=args.aux, se_loss=args.se_loss, norm_layer=BatchNorm, 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'], strict=False) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) # print(model) scales = [0.5, 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] if not args.ms: scales = [1.0] evaluator = MultiEvalModule(model, testset.num_class, scales=scales, flip=args.ms).cuda() evaluator.eval() metric = utils.SegmentationMetric(testset.num_class) tbar = tqdm(test_data) total_inter, total_union, total_correct, total_label, all_label = 0, 0, 0, 0, 0 # for i, (image, dst) in enumerate(tbar): # # print(dst) # with torch.no_grad(): # outputs = evaluator.parallel_forward(image)[0] # correct, labeled = batch_pix_accuracy(outputs, dst[0]) # total_correct += correct # all_label += labeled # img_pixAcc = 1.0 * correct / (np.spacing(1) + labeled) # inter, union, area_pred, area_lab = batch_intersection_union(outputs, dst[0], testset.num_class) # total_label += area_lab # total_inter += inter # total_union += union # class_pixAcc = 1.0 * inter / (np.spacing(1) + area_lab) # class_IoU = 1.0 * inter / (np.spacing(1) + union) # class_mIoU = class_IoU.mean() # print("img pixAcc:", img_pixAcc) # print("img Classes pixAcc:", class_pixAcc) # print("img Classes IoU:", class_IoU) # total_pixAcc = 1.0 * total_correct / (np.spacing(1) + all_label) # pixAcc = 1.0 * total_inter / (np.spacing(1) + total_label) # IoU = 1.0 * total_inter / (np.spacing(1) + total_union) # mIoU = IoU.mean() # print("set pixAcc:", pixAcc) # print("set Classes pixAcc:", pixAcc) # print("set Classes IoU:", IoU) # print("set mean IoU:", mIoU) for i, (image, dst) in enumerate(tbar): if 'val' in args.mode: with torch.no_grad(): predicts = evaluator.parallel_forward(image) # metric.update(dst[0], predicts[0]) # pixAcc, mIoU = metric.get() # tbar.set_description( 'pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU)) else: with torch.no_grad(): outputs = evaluator.parallel_forward(image)