segmentation/tool/train.py [63:141]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - assert (args.train_h - 1) % 8 == 0 and (args.train_w - 1) % 8 == 0 elif args.arch == 'psa': if args.compact: args.mask_h = (args.train_h - 1) // (8 * args.shrink_factor) + 1 args.mask_w = (args.train_w - 1) // (8 * args.shrink_factor) + 1 else: assert (args.mask_h is None and args.mask_w is None) or ( args.mask_h is not None and args.mask_w is not None) if args.mask_h is None and args.mask_w is None: args.mask_h = 2 * ((args.train_h - 1) // (8 * args.shrink_factor) + 1) - 1 args.mask_w = 2 * ((args.train_w - 1) // (8 * args.shrink_factor) + 1) - 1 else: assert (args.mask_h % 2 == 1) and (args.mask_h >= 3) and ( args.mask_h <= 2 * ((args.train_h - 1) // (8 * args.shrink_factor) + 1) - 1) assert (args.mask_w % 2 == 1) and (args.mask_w >= 3) and ( args.mask_w <= 2 * ((args.train_h - 1) // (8 * args.shrink_factor) + 1) - 1) else: raise Exception('architecture not supported yet'.format(args.arch)) def main(): args = get_parser() check(args) os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.train_gpu) if args.manual_seed is not None: random.seed(args.manual_seed) np.random.seed(args.manual_seed) torch.manual_seed(args.manual_seed) torch.cuda.manual_seed(args.manual_seed) torch.cuda.manual_seed_all(args.manual_seed) cudnn.benchmark = False cudnn.deterministic = True if args.dist_url == "env://" and args.world_size == -1: args.world_size = int(os.environ["WORLD_SIZE"]) args.distributed = args.world_size > 1 or args.multiprocessing_distributed args.ngpus_per_node = len(args.train_gpu) if len(args.train_gpu) == 1: args.sync_bn = False args.distributed = False args.multiprocessing_distributed = False else: args.distributed = False if args.multiprocessing_distributed: port = find_free_port() args.dist_url = f"tcp://127.0.0.1:{port}" args.world_size = args.ngpus_per_node * args.world_size mp.spawn(main_worker, nprocs=args.ngpus_per_node, args=(args.ngpus_per_node, args)) else: main_worker(args.train_gpu, args.ngpus_per_node, args) def main_worker(gpu, ngpus_per_node, argss): global args args = argss if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) criterion = nn.CrossEntropyLoss(ignore_index=args.ignore_label) if args.arch == 'psp': from model.pspnet import PSPNet model = PSPNet(layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, criterion=criterion) modules_ori = [model.layer0, model.layer1, model.layer2, model.layer3, model.layer4] modules_new = [model.ppm, model.cls, model.aux] elif args.arch == 'psa': from model.psanet import PSANet model = PSANet(layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, psa_type=args.psa_type, compact=args.compact, shrink_factor=args.shrink_factor, mask_h=args.mask_h, mask_w=args.mask_w, normalization_factor=args.normalization_factor, psa_softmax=args.psa_softmax, criterion=criterion) modules_ori = [model.layer0, model.layer1, model.layer2, model.layer3, model.layer4] modules_new = [model.psa, model.cls, model.aux] elif args.arch == 'fcn': from model.fcn import FCNet model = FCNet(layers=args.layers, classes=args.classes, criterion=criterion) modules_ori = [model.model.backbone] modules_new = [model.model.classifier, model.model.aux_classifier] - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - segmentation/tool/train_cnsn.py [63:141]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - assert (args.train_h - 1) % 8 == 0 and (args.train_w - 1) % 8 == 0 elif args.arch == 'psa': if args.compact: args.mask_h = (args.train_h - 1) // (8 * args.shrink_factor) + 1 args.mask_w = (args.train_w - 1) // (8 * args.shrink_factor) + 1 else: assert (args.mask_h is None and args.mask_w is None) or ( args.mask_h is not None and args.mask_w is not None) if args.mask_h is None and args.mask_w is None: args.mask_h = 2 * ((args.train_h - 1) // (8 * args.shrink_factor) + 1) - 1 args.mask_w = 2 * ((args.train_w - 1) // (8 * args.shrink_factor) + 1) - 1 else: assert (args.mask_h % 2 == 1) and (args.mask_h >= 3) and ( args.mask_h <= 2 * ((args.train_h - 1) // (8 * args.shrink_factor) + 1) - 1) assert (args.mask_w % 2 == 1) and (args.mask_w >= 3) and ( args.mask_w <= 2 * ((args.train_h - 1) // (8 * args.shrink_factor) + 1) - 1) else: raise Exception('architecture not supported yet'.format(args.arch)) def main(): args = get_parser() check(args) os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.train_gpu) if args.manual_seed is not None: random.seed(args.manual_seed) np.random.seed(args.manual_seed) torch.manual_seed(args.manual_seed) torch.cuda.manual_seed(args.manual_seed) torch.cuda.manual_seed_all(args.manual_seed) cudnn.benchmark = False cudnn.deterministic = True if args.dist_url == "env://" and args.world_size == -1: args.world_size = int(os.environ["WORLD_SIZE"]) args.distributed = args.world_size > 1 or args.multiprocessing_distributed args.ngpus_per_node = len(args.train_gpu) if len(args.train_gpu) == 1: args.sync_bn = False args.distributed = False args.multiprocessing_distributed = False else: args.distributed = False if args.multiprocessing_distributed: port = find_free_port() args.dist_url = f"tcp://127.0.0.1:{port}" args.world_size = args.ngpus_per_node * args.world_size mp.spawn(main_worker, nprocs=args.ngpus_per_node, args=(args.ngpus_per_node, args)) else: main_worker(args.train_gpu, args.ngpus_per_node, args) def main_worker(gpu, ngpus_per_node, argss): global args args = argss if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) criterion = nn.CrossEntropyLoss(ignore_index=args.ignore_label) if args.arch == 'psp': from model.pspnet import PSPNet model = PSPNet(layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, criterion=criterion) modules_ori = [model.layer0, model.layer1, model.layer2, model.layer3, model.layer4] modules_new = [model.ppm, model.cls, model.aux] elif args.arch == 'psa': from model.psanet import PSANet model = PSANet(layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, psa_type=args.psa_type, compact=args.compact, shrink_factor=args.shrink_factor, mask_h=args.mask_h, mask_w=args.mask_w, normalization_factor=args.normalization_factor, psa_softmax=args.psa_softmax, criterion=criterion) modules_ori = [model.layer0, model.layer1, model.layer2, model.layer3, model.layer4] modules_new = [model.psa, model.cls, model.aux] elif args.arch == 'fcn': from model.fcn import FCNet model = FCNet(layers=args.layers, classes=args.classes, criterion=criterion) modules_ori = [model.model.backbone] modules_new = [model.model.classifier, model.model.aux_classifier] - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -