if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('-e', '--epochs', default='last', type=str)
parser.add_argument('-d', '--devices', default='1', type=str)
parser.add_argument('-v', '--verbose', default=False, action='store_true')
parser.add_argument('--show_image',
'-s',
default=False,
action='store_true')
parser.add_argument('--save_path', '-p', default=None)
args = parser.parse_args()
all_dev = parse_devices(args.devices)
mp_ctx = mp.get_context('spawn')
network = CPNet(config.num_classes, criterion=None)
data_setting = {
'img_root': config.img_root_folder,
'gt_root': config.gt_root_folder,
'train_source': config.train_source,
'eval_source': config.eval_source
}
dataset = ADE(data_setting, 'val', None)
with torch.no_grad():
segmentor = SegEvaluator(dataset, config.num_classes,
config.image_mean, config.image_std, network,
config.eval_scale_array, config.eval_flip,
all_dev, args.verbose, args.save_path,
args.show_image)
segmentor.run(config.snapshot_dir, args.epochs, config.val_log_file,
if engine.distributed:
torch.cuda.set_device(engine.local_rank)
# data loader
train_loader, train_sampler = get_train_loader(engine, ADE)
# config network and criterion
criterion = nn.CrossEntropyLoss(reduction='mean', ignore_index=-1)
if engine.distributed:
logger.info('Use the Multi-Process-SyncBatchNorm')
BatchNorm2d = SyncBatchNorm
# else:
# BatchNorm2d = BatchNorm2d
model = CPNet(config.num_classes,
criterion=criterion,
pretrained_model=config.pretrained_model,
norm_layer=BatchNorm2d)
init_weight(model.business_layer,
nn.init.kaiming_normal_,
BatchNorm2d,
config.bn_eps,
config.bn_momentum,
mode='fan_in',
nonlinearity='relu')
# group weight and config optimizer
base_lr = config.lr
# if engine.distributed:
# base_lr = config.lr * engine.world_size
params_list = []