├── README.md ├── exps └── SemanticSceneCompletion │ ├── config.py │ ├── dataloader.py │ ├── eval.py │ ├── network.py │ ├── nyu.py │ ├── st-run.sh │ ├── test.sh │ └── train.py ├── figs ├── nyu_res.png └── nyu_vis.png └── furnace ├── .DS_Store ├── __init__.py ├── base_model ├── README.md ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ └── resnet.cpython-36.pyc ├── resnet.py └── xception.py ├── datasets ├── .DS_Store ├── BaseDataset.py ├── __init__.py └── __pycache__ │ ├── BaseDataset.cpython-36.pyc │ ├── BaseDataset.cpython-37.pyc │ ├── BaseDataset.cpython-38.pyc │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ └── __init__.cpython-38.pyc ├── engine ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── engine.cpython-36.pyc │ ├── engine.cpython-37.pyc │ ├── engine.cpython-38.pyc │ ├── evaluator.cpython-37.pyc │ ├── evaluator.cpython-38.pyc │ ├── logger.cpython-36.pyc │ ├── logger.cpython-37.pyc │ ├── logger.cpython-38.pyc │ ├── lr_policy.cpython-36.pyc │ ├── lr_policy.cpython-37.pyc │ ├── lr_policy.cpython-38.pyc │ ├── version.cpython-36.pyc │ ├── version.cpython-37.pyc │ └── version.cpython-38.pyc ├── engine.py ├── evaluator.py ├── logger.py ├── lr_policy.py └── version.py ├── seg_opr ├── .DS_Store ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── metric.cpython-37.pyc │ ├── metric.cpython-38.pyc │ ├── seg_oprs.cpython-36.pyc │ ├── seg_oprs.cpython-37.pyc │ └── seg_oprs.cpython-38.pyc ├── metric.py ├── parallel │ └── parallel_apply.py ├── seg_oprs.py ├── sgd.py └── sync_bn │ ├── __init__.py │ ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── comm.cpython-36.pyc │ ├── comm.cpython-37.pyc │ ├── comm.cpython-38.pyc │ ├── functions.cpython-36.pyc │ ├── functions.cpython-37.pyc │ ├── functions.cpython-38.pyc │ ├── parallel.cpython-36.pyc │ ├── parallel.cpython-37.pyc │ ├── parallel.cpython-38.pyc │ ├── syncbn.cpython-36.pyc │ ├── syncbn.cpython-37.pyc │ └── syncbn.cpython-38.pyc │ ├── comm.py │ ├── functions.py │ ├── parallel.py │ ├── parallel_apply.py │ ├── src │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── __init__.cpython-37.pyc │ │ └── __init__.cpython-38.pyc │ ├── cpu │ │ ├── .ninja_deps │ │ ├── .ninja_log │ │ ├── __init__.py │ │ ├── build.ninja │ │ ├── dist │ │ │ └── syncbn_cpu-0.0.0-py3.6-linux-x86_64.egg │ │ ├── operator.cpp │ │ ├── operator.h │ │ ├── operator.o │ │ ├── setup.py │ │ ├── syncbn_cpu.cpp │ │ ├── syncbn_cpu.egg-info │ │ │ └── PKG-INFO │ │ ├── syncbn_cpu.o │ │ └── syncbn_cpu.so │ └── gpu │ │ ├── .ninja_deps │ │ ├── .ninja_log │ │ ├── __init__.py │ │ ├── build.ninja │ │ ├── common.h │ │ ├── device_tensor.h │ │ ├── dist │ │ └── syncbn_gpu-0.0.0-py3.6-linux-x86_64.egg │ │ ├── operator.cpp │ │ ├── operator.h │ │ ├── operator.o │ │ ├── setup.py │ │ ├── syncbn_gpu.egg-info │ │ └── PKG-INFO │ │ ├── syncbn_gpu.so │ │ ├── syncbn_kernel.cu │ │ └── syncbn_kernel.cuda.o │ └── syncbn.py └── utils ├── __init__.py ├── __pycache__ ├── __init__.cpython-36.pyc ├── __init__.cpython-37.pyc ├── __init__.cpython-38.pyc ├── img_utils.cpython-36.pyc ├── img_utils.cpython-37.pyc ├── img_utils.cpython-38.pyc ├── init_func.cpython-36.pyc ├── init_func.cpython-37.pyc ├── init_func.cpython-38.pyc ├── pyt_utils.cpython-36.pyc ├── pyt_utils.cpython-37.pyc ├── pyt_utils.cpython-38.pyc ├── visualize.cpython-37.pyc └── visualize.cpython-38.pyc ├── img_utils.py ├── init_func.py ├── pyt_utils.py └── visualize.py /README.md: -------------------------------------------------------------------------------- 1 | # SISNet 2 | Semantic Scene Completion via Integrating Instances and Scene in-the-Loop (CVPR 2021) 3 | 4 | In this repository, we provide SISNet model implementation (with Pytorch) as well as data preparation, training and evaluation scripts on NYU Depth V2, NYUCAD and SUNCG_RGBD. 5 | 6 | ![image](https://github.com/yjcaimeow/SISNet/blob/main/figs/nyu_vis.png) 7 | 8 | ##### Code is under construction. 9 | 10 | ## Getting Started 11 | ### Set up 12 | Clone the repository: 13 | 14 | git clone https://github.com/yjcaimeow/SISNet 15 | ### Installation 16 | The code is tested with Ubuntu 18.04, Pytorch v1.4.0, CUDA 10.0 and cuDNN v7.4. 17 | 18 | conda install pytorch==1.4.0 torchvision==0.5.0 cudatoolkit=10.0 -c pytorch 19 | Install the following Python dependencies (with pip install ): 20 | 21 | matplotlib 22 | opencv-python 23 | plyfile 24 | 'trimesh>=2.35.39,<2.35.40' 25 | 'networkx>=2.2,<2.3' 26 | tqdm 27 | ninja 28 | easydict 29 | argparse 30 | h5py 31 | scipy 32 | Compile the CUDA layers , which we used in the network: 33 | 34 | cd pointnet2 35 | python setup.py install --user 36 | 37 | # Chamfer Distancecd 38 | cd extensions/chamfer_dist 39 | python setup.py install --user 40 | 41 | # Cubic Feature Samplingcd 42 | cd extensions/cubic_feature_sampling 43 | python setup.py install --user 44 | 45 | # Gridding & Gridding Reversecd 46 | cd extensions/gridding 47 | python setup.py install --user 48 | 49 | # apex for multi-gpus train 50 | cd ./furnace/apex 51 | export CUDA_HOME=/usr/local/cuda-10.0 52 | pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./ 53 | 54 | ### Datasets 55 | We use the nyu depth v2, nyucad and suncg_rgbd datasets in our experiments, which are available in ./data folder. The google drive link provides the input data we need (nyu depth v2 and nyucad). 56 | 57 | Note that, for suncg_rgbd we follow and obtrain the data from SATNet. https://github.com/ShiceLiu/SATNet 58 | 59 | For suncg_rgbd images, please download from the following link and unzip the SUNCGRGBD_images.zip under data/suncg/ https://pan.baidu.com/s/1vUt09GlkC1lPFRm8zXofZA 60 | 61 | For the pre-process 2D semantic segmentation, we just use the following semantic 62 | segmentation code (https://github.com/Tramac/awesome-semantic-segmentation-pytorch) which includes many classical methods. So you can train your segmention to obtain the input data of initial scene completion or directly use the npz files provided in the Google Drive (voxel_bisenet_res34_4e-3). 63 | 64 | ### Get Started 65 | To train SISNet, you can simply use the following command: 66 | 67 | bash scripts/train_SISNet.sh 68 | 69 | To test SISNet, you can use the following command: 70 | 71 | bash scripts/test_SISNet.sh 72 | 73 | The results are save in results/${dataset}. You can use scripts/toply.py to transform the 74 | npz result to ply and use MeshLab to show. 75 | 76 | Finally, the project folder is organized as follows. 77 | 78 | SISNet 79 | ├── core 80 | │ ├── scene_completion 81 | │ ├── instance_completion 82 | ├── data 83 | │ ├── nyucad 84 | │ ├── suncg 85 | │ ├── nyuv2 86 | │ │ ├── train.txt 87 | │ │ ├── test.txt 88 | │ │ ├── R （RGB images） 89 | │ │ ├── L (Label ground truth of semantic scene completion) 90 | │ │ ├── M (Mapping from 2D to 3D) 91 | │ │ ├── T (TSDF truncated signed distance function) 92 | │ │ ├── voxel_bisenet_res34_4e-3（sem of visiable region） 93 | ├── instance_data 94 | │ ├── nyuv2 95 | ├── nyuv2_S*_test/train_result 96 | ├── ckpts 97 | │ ├── nyuv2_s*.pth 98 | ├── I*_ckpt 99 | ├── results 100 | │ ├── nyuv2 101 | │ │ ├── *_test/train_results 102 | ├── utils 103 | │ ├── vis_ssc.py 104 | │ ├── ... 105 | 106 | ### Results on NYU Depth v2 Benchmark 107 | ![image](https://github.com/yjcaimeow/SISNet/blob/main/figs/nyu_res.png) 108 | -------------------------------------------------------------------------------- /exps/SemanticSceneCompletion/config.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | 3 | from __future__ import absolute_import 4 | from __future__ import division 5 | from __future__ import print_function 6 | 7 | import os 8 | import os.path as osp 9 | import sys 10 | import time 11 | import numpy as np 12 | from easydict import EasyDict as edict 13 | import argparse 14 | 15 | C = edict() 16 | config = C 17 | cfg = C 18 | 19 | C.seed = 12345 20 | 21 | remoteip = os.popen('pwd').read() 22 | 23 | C.volna = '/mnt/lustre/caiyingjie/' 24 | 25 | """please config ROOT_dir and user when u first using""" 26 | C.abs_dir = osp.realpath(".") 27 | C.this_dir = C.abs_dir.split(osp.sep)[-1] 28 | 29 | 30 | C.root_dir = '/mnt/lustre/liushinan/cyj/start' 31 | C.log_dir = osp.abspath('logs') 32 | C.tb_dir = osp.abspath(osp.join(C.log_dir, "tb")) 33 | 34 | C.log_dir_link = C.log_dir 35 | C.snapshot_dir = osp.abspath(osp.join(C.log_dir, "models")) 36 | 37 | exp_time = time.strftime('%Y_%m_%d_%H_%M_%S', time.localtime()) 38 | C.log_file = C.log_dir + '/log_' + exp_time + '.log' 39 | C.link_log_file = C.log_file + '/log_last.log' 40 | C.val_log_file = C.log_dir + '/_' + exp_time + '.lovalg' 41 | C.link_val_log_file = C.log_dir + '/val_last.log' 42 | 43 | """Data Dir and Weight Dir""" 44 | C.dataset_path = './data' 45 | C.i_root_folder = C.dataset_path 46 | C.g_root_folder = C.dataset_path 47 | C.h_root_folder = osp.join(C.dataset_path, 'H') 48 | C.m_root_folder = osp.join(C.dataset_path, 'M') 49 | C.t_source = osp.join(C.dataset_path, "train.txt") 50 | C.e_source = osp.join(C.dataset_path, "test.txt") 51 | C.is_test = False 52 | 53 | """Path Config""" 54 | def add_path(path): 55 | if path not in sys.path: 56 | sys.path.insert(0, path) 57 | 58 | 59 | add_path(osp.join(C.root_dir, 'furnace')) 60 | 61 | """Image Config""" 62 | C.num_classes = 12 63 | C.background = 255 64 | C.image_mean = np.array([0.485, 0.456, 0.406]) # 0.485, 0.456, 0.406 65 | C.image_std = np.array([0.229, 0.224, 0.225]) 66 | C.num_train_imgs = 795 67 | C.num_eval_imgs = 815 68 | 69 | """ Settings for network, this would be different for each kind of model""" 70 | C.fix_bias = True 71 | C.bn_eps = 1e-5 72 | C.bn_momentum = 0.1 73 | #C.pretrained_model = C.volna + 'DATA/.torch/pytorch-weight/resnet50-imagenet.pth' 74 | C.aux_loss_alpha = 0.1 75 | 76 | """Train Config""" 77 | C.lr = 0.2 78 | C.lr_power = 0.9 79 | C.momentum = 0.9 80 | C.weight_decay = 5e-4 81 | C.batch_size = 64 82 | C.nepochs = 250 83 | C.niters_per_epoch = 795 // C.batch_size 84 | 85 | C.num_workers = 0 86 | 87 | C.train_scale_array = [1] 88 | C.warm_up_epoch = 0 89 | 90 | """Eval Config""" 91 | C.eval_iter = 30 92 | C.eval_stride_rate = 2 / 3 93 | C.eval_scale_array = [1, ] 94 | C.eval_flip = False 95 | C.eval_base_size = 480 96 | C.eval_crop_size = 640 97 | 98 | """Display Config""" 99 | C.snapshot_iter = 10 100 | C.record_info_iter = 20 101 | C.display_iter = 50 102 | C.edge_weight = 1 103 | C.edge_weight_gsnn = 1.5 104 | 105 | C.kld_weight = 2 106 | C.samples = 4 107 | C.lantent_size = 16 108 | 109 | def open_tensorboard(): 110 | pass 111 | 112 | if __name__ == '__main__': 113 | print(config.nepochs) 114 | parser = argparse.ArgumentParser() 115 | parser.add_argument( 116 | '-tb', '--tensorboard', default=True, action='store_true') 117 | args = parser.parse_args() 118 | 119 | if args.tensorboard: 120 | open_tensorboard() 121 | -------------------------------------------------------------------------------- /exps/SemanticSceneCompletion/dataloader.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | import cv2 3 | import torch 4 | import numpy as np 5 | from torch.utils import data 6 | import random 7 | from config import config 8 | import sys 9 | sys.path.append('../../furnace/') 10 | from utils.img_utils import normalize, \ 11 | generate_random_crop_pos, random_crop_pad_to_shape 12 | 13 | class TrainPre(object): 14 | def __init__(self, img_mean, img_std): 15 | self.img_mean = img_mean 16 | self.img_std = img_std 17 | 18 | def __call__(self, img): 19 | img = normalize(img, self.img_mean, self.img_std) 20 | 21 | p_img = img.transpose(2, 0, 1) 22 | extra_dict = {} 23 | 24 | return p_img, extra_dict 25 | class ValPre(object): 26 | def __call__(self, img): 27 | extra_dict = {} 28 | return img, extra_dict 29 | 30 | 31 | def get_train_loader(engine, dataset, s3client): 32 | # ''' 33 | data_setting = {'i_root': config.i_root_folder, 34 | 'g_root': config.g_root_folder, 35 | 'h_root':config.h_root_folder, 36 | 'ROOT_DIR': config.dataset_path, 37 | 'LOG_DIR': config.log_dir, 38 | 'm_root': config.m_root_folder, 39 | 't_source': config.t_source, 40 | 'e_source': config.e_source} 41 | 42 | train_preprocess = TrainPre(config.image_mean, config.image_std) 43 | 44 | train_dataset = dataset(data_setting, "train", 45 | train_preprocess, config.batch_size * config.niters_per_epoch, s3client=s3client) 46 | 47 | train_sampler = None 48 | is_shuffle = True 49 | batch_size = config.batch_size 50 | 51 | if engine.distributed: 52 | train_sampler = torch.utils.data.distributed.DistributedSampler( 53 | train_dataset) 54 | batch_size = config.batch_size // engine.world_size 55 | is_shuffle = False 56 | train_loader = data.DataLoader(train_dataset, 57 | batch_size=batch_size, 58 | num_workers=config.num_workers, 59 | drop_last=False, 60 | shuffle=is_shuffle, 61 | pin_memory=True, 62 | sampler=train_sampler) 63 | 64 | return train_loader, train_sampler 65 | -------------------------------------------------------------------------------- /exps/SemanticSceneCompletion/eval.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | import os 4 | import cv2 5 | import argparse 6 | import numpy as np 7 | 8 | import torch 9 | import torch.nn as nn 10 | import torch.multiprocessing as mp 11 | 12 | from config import config 13 | import sys 14 | #sys.path.append('/mnt/lustre/liushinan/cyj/start/furnace') 15 | from utils.pyt_utils import ensure_dir, link_file, load_model, parse_devices 16 | from utils.visualize import print_iou, show_img 17 | from engine.evaluator import Evaluator 18 | from engine.logger import get_logger 19 | from seg_opr.metric import hist_info, compute_score 20 | from nyu import NYUv2 21 | from network import Network 22 | from dataloader import ValPre 23 | 24 | logger = get_logger() 25 | 26 | 27 | class SegEvaluator(Evaluator): 28 | def func_per_iteration(self, data, device): 29 | lll = data['lll'] 30 | www = data['www'] 31 | mmm = data['mmm'] 32 | tsdf = data['tsdf'] 33 | ttt = data['ttt'] 34 | 35 | name = data['fn'] 36 | pp, ps, pe = self.eval_ssc(ttt, tsdf, device) 37 | 38 | results_dict = {'pp':pp, 'll':lll, 'ww':www, 'pred': ps, 'tsdf':tsdf, 39 | 'name':name, 'mm':mmm} 40 | return results_dict 41 | 42 | def hist_info(self, n_cl, ppp, ggg): 43 | assert (ppp.shape == ppp.shape) 44 | k = (ggg >= 0) & (ggg < n_cl) 45 | labeled = np.sum(k) 46 | correct = np.sum((ppp[k] == ggg[k])) 47 | 48 | return np.bincount(n_cl * ggg[k].astype(int) + ppp[k].astype(int), 49 | minlength=n_cl ** 2).reshape(n_cl, 50 | n_cl), correct, labeled 51 | 52 | def compute_metric(self, results): 53 | hist_ssc = np.zeros((config.num_classes, config.num_classes)) 54 | correct_ssc = 0 55 | labeled_ssc = 0 56 | 57 | tp_sc, fp_sc, fn_sc, union_sc, intersection_sc = 0, 0, 0, 0, 0 58 | for d in results: 59 | ppp = d['pp'].astype(np.int64) 60 | lll = d['ll'].astype(np.int64) 61 | www = d['ww'].astype(np.float32) 62 | mmm = d['mm'].astype(np.int64).reshape(-1) 63 | name= d['name'] 64 | flat_ppp = np.ravel(ppp) 65 | flat_lll = np.ravel(lll) 66 | tsdf = np.ravel(d['tsdf']) 67 | 68 | nff = np.where((www > 0))# & (tsdf>0)) 69 | nff_ppp = flat_ppp[nff] 70 | nff_lll = flat_lll[nff] 71 | 72 | h_ssc, c_ssc, l_ssc = self.hist_info(config.num_classes, nff_ppp, nff_lll) 73 | hist_ssc += h_ssc 74 | correct_ssc += c_ssc 75 | labeled_ssc += l_ssc 76 | 77 | ooo = (mmm == 307200) & (www > 0) & (flat_lll != 255) #& (tsdf>0) 78 | ooo_ppp = flat_ppp[ooo] 79 | ooo_lll = flat_lll[ooo] 80 | 81 | tp_occ = ((ooo_lll > 0) & (ooo_ppp > 0)).astype(np.int8).sum() 82 | fp_occ = ((ooo_lll == 0) & (ooo_ppp > 0)).astype(np.int8).sum() 83 | fn_occ = ((ooo_lll > 0) & (ooo_ppp == 0)).astype(np.int8).sum() 84 | 85 | union = ((ooo_lll > 0) | (ooo_ppp > 0)).astype(np.int8).sum() 86 | intersection = ((ooo_lll > 0) & (ooo_ppp > 0)).astype(np.int8).sum() 87 | 88 | tp_sc += tp_occ 89 | fp_sc += fp_occ 90 | fn_sc += fn_occ 91 | union_sc += union 92 | intersection_sc += intersection 93 | 94 | score_ssc = compute_score(hist_ssc, correct_ssc, labeled_ssc) 95 | IOU_sc = intersection_sc / union_sc 96 | precision_sc = tp_sc / (tp_sc + fp_sc) 97 | recall_sc = tp_sc / (tp_sc + fn_sc) 98 | score_sc = [IOU_sc, precision_sc, recall_sc] 99 | 100 | result_line = self.print_ssc_iou(score_sc, score_ssc) 101 | return result_line 102 | 103 | def eval_ssc(self, ttt, tsdf, device=None): 104 | sc, bsc, esc = self.val_func_process_ssc(ttt, tsdf, device) 105 | sc = sc.permute(1, 2, 3, 0) ## 60 36 60 12 106 | softmax = nn.Softmax(dim=3) 107 | sc = softmax(sc) 108 | ddddd = sc.cpu().numpy() 109 | pp = ddddd.argmax(3) 110 | 111 | return pp, ddddd, None 112 | 113 | def val_func_process_ssc(self, ttt, tsdf, device=None): 114 | 115 | tsdf = np.ascontiguousarray(tsdf[None, :], dtype=np.float32) 116 | tsdf = torch.FloatTensor(tsdf).cuda(device) 117 | 118 | ttt = np.ascontiguousarray(ttt[None, :], dtype=np.float32) 119 | ttt = torch.FloatTensor(ttt).cuda(device) 120 | with torch.cuda.device(ttt.get_device()): 121 | self.val_func.eval() 122 | self.val_func.to(ttt.get_device()) 123 | with torch.no_grad(): 124 | sc, ssc = self.val_func(ttt, tsdf) 125 | sc = sc[0] 126 | sc = torch.exp(sc) 127 | return sc, ssc, ssc 128 | 129 | def print_ssc_iou(self, sc, ssc): 130 | lines = [] 131 | lines.append('--*-- Semantic Scene Completion --*--') 132 | lines.append('IOU: \n{}\n'.format(str(ssc[0].tolist()))) 133 | lines.append('meanIOU: %f\n' % ssc[2]) 134 | lines.append('pixel-accuracy: %f\n' % ssc[3]) 135 | lines.append('') 136 | lines.append('--*-- Scene Completion --*--\n') 137 | lines.append('IOU: %f\n' % sc[0]) 138 | lines.append('pixel-accuracy: %f\n' % sc[1]) 139 | lines.append('recall: %f\n' % sc[2]) 140 | 141 | line = "\n".join(lines) 142 | print(line) 143 | return line 144 | 145 | if __name__ == "__main__": 146 | parser = argparse.ArgumentParser() 147 | parser.add_argument('-e', '--epochs', default='last', type=str) 148 | parser.add_argument('-d', '--devices', default='1', type=str) 149 | parser.add_argument('-v', '--verbose', default=False, action='store_true') 150 | 151 | args = parser.parse_args() 152 | all_dev = parse_devices(args.devices) 153 | 154 | network = Network(class_num=config.num_classes, feature=128, bn_momentum=config.bn_momentum, 155 | norm_layer=nn.BatchNorm3d, eval=True) 156 | data_setting = {'i_root': config.i_root_folder, 157 | 'g_root': config.g_root_folder, 158 | 'h_root':config.h_root_folder, 159 | 'm_root': config.m_root_folder, 160 | 't_source': config.t_source, 161 | 'e_source': config.e_source} 162 | val_pre = ValPre() 163 | dataset = NYUv2(data_setting, 'val', val_pre) 164 | 165 | with torch.no_grad(): 166 | segmentor = SegEvaluator(dataset, config.num_classes, config.image_mean, 167 | config.image_std, network, 168 | config.eval_scale_array, config.eval_flip, 169 | all_dev, args.verbose) 170 | segmentor.run(config.snapshot_dir, args.epochs, config.val_log_file, 171 | config.link_val_log_file) 172 | -------------------------------------------------------------------------------- /exps/SemanticSceneCompletion/network.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | from functools import partial 7 | from config import config 8 | 9 | class SimpleRB(nn.Module): 10 | def __init__(self, in_channel, norm_layer, bn_momentum): 11 | super(SimpleRB, self).__init__() 12 | self.path = nn.Sequential( 13 | nn.Conv3d(in_channel, in_channel, kernel_size=3, padding=1, bias=False), 14 | norm_layer(in_channel, momentum=bn_momentum), 15 | nn.ReLU(), 16 | nn.Conv3d(in_channel, in_channel, kernel_size=3, padding=1, bias=False), 17 | norm_layer(in_channel, momentum=bn_momentum), 18 | ) 19 | self.relu = nn.ReLU() 20 | 21 | def forward(self, x): 22 | residual = x 23 | conv_path = self.path(x) 24 | out = residual + conv_path 25 | out = self.relu(out) 26 | return out 27 | 28 | class Bottleneck3D(nn.Module): 29 | 30 | def __init__(self, inplanes, planes, norm_layer, stride=1, dilation=[1, 1, 1], expansion=4, downsample=None, 31 | fist_dilation=1, multi_grid=1, 32 | bn_momentum=0.0003): 33 | super(Bottleneck3D, self).__init__() 34 | self.expansion = expansion 35 | self.conv1 = nn.Conv3d(inplanes, planes, kernel_size=1, bias=False) 36 | self.bn1 = norm_layer(planes, momentum=bn_momentum) 37 | self.conv2 = nn.Conv3d(planes, planes, kernel_size=(1, 1, 3), stride=(1, 1, stride), 38 | dilation=(1, 1, dilation[0]), padding=(0, 0, dilation[0]), bias=False) 39 | self.bn2 = norm_layer(planes, momentum=bn_momentum) 40 | self.conv3 = nn.Conv3d(planes, planes, kernel_size=(1, 3, 1), stride=(1, stride, 1), 41 | dilation=(1, dilation[1], 1), padding=(0, dilation[1], 0), bias=False) 42 | self.bn3 = norm_layer(planes, momentum=bn_momentum) 43 | self.conv4 = nn.Conv3d(planes, planes, kernel_size=(3, 1, 1), stride=(stride, 1, 1), 44 | dilation=(dilation[2], 1, 1), padding=(dilation[2], 0, 0), bias=False) 45 | self.bn4 = norm_layer(planes, momentum=bn_momentum) 46 | self.conv5 = nn.Conv3d(planes, planes * self.expansion, kernel_size=(1, 1, 1), bias=False) 47 | self.bn5 = norm_layer(planes * self.expansion, momentum=bn_momentum) 48 | 49 | self.relu = nn.ReLU(inplace=False) 50 | self.relu_inplace = nn.ReLU(inplace=True) 51 | self.downsample = downsample 52 | self.dilation = dilation 53 | self.stride = stride 54 | 55 | self.downsample2 = nn.Sequential( 56 | nn.AvgPool3d(kernel_size=(1, stride, 1), stride=(1, stride, 1)), 57 | nn.Conv3d(planes, planes, kernel_size=1, stride=1, bias=False), 58 | norm_layer(planes, momentum=bn_momentum), 59 | ) 60 | self.downsample3 = nn.Sequential( 61 | nn.AvgPool3d(kernel_size=(stride, 1, 1), stride=(stride, 1, 1)), 62 | nn.Conv3d(planes, planes, kernel_size=1, stride=1, bias=False), 63 | norm_layer(planes, momentum=bn_momentum), 64 | ) 65 | self.downsample4 = nn.Sequential( 66 | nn.AvgPool3d(kernel_size=(stride, 1, 1), stride=(stride, 1, 1)), 67 | nn.Conv3d(planes, planes, kernel_size=1, stride=1, bias=False), 68 | norm_layer(planes, momentum=bn_momentum), 69 | ) 70 | 71 | def forward(self, x): 72 | residual = x 73 | 74 | out1 = self.relu(self.bn1(self.conv1(x))) 75 | out2 = self.bn2(self.conv2(out1)) 76 | out2_relu = self.relu(out2) 77 | 78 | out3 = self.bn3(self.conv3(out2_relu)) 79 | if self.stride != 1: 80 | out2 = self.downsample2(out2) 81 | out3 = out3 + out2 82 | out3_relu = self.relu(out3) 83 | 84 | out4 = self.bn4(self.conv4(out3_relu)) 85 | if self.stride != 1: 86 | out2 = self.downsample3(out2) 87 | out3 = self.downsample4(out3) 88 | out4 = out4 + out2 + out3 89 | 90 | out4_relu = self.relu(out4) 91 | out5 = self.bn5(self.conv5(out4_relu)) 92 | 93 | if self.downsample is not None: 94 | residual = self.downsample(x) 95 | 96 | out = out5 + residual 97 | out_relu = self.relu(out) 98 | 99 | return out_relu 100 | 101 | class STAGE2(nn.Module): 102 | def __init__(self, class_num, norm_layer, resnet_out=2048, feature=512, ThreeDinit=True, 103 | bn_momentum=0.1, pretrained_model=None, eval=False, freeze_bn=False): 104 | super(STAGE2, self).__init__() 105 | self.business_layer = [] 106 | if eval: 107 | self.downsample = nn.Sequential( 108 | nn.Conv2d(resnet_out, feature, kernel_size=1, bias=False), 109 | nn.BatchNorm2d(feature, momentum=bn_momentum), 110 | nn.ReLU() 111 | ) 112 | else: 113 | self.downsample = nn.Sequential( 114 | nn.Conv2d(resnet_out, feature, kernel_size=1, bias=False), 115 | norm_layer(feature, momentum=bn_momentum), 116 | nn.ReLU() 117 | ) 118 | self.business_layer.append(self.downsample) 119 | self.resnet_out = resnet_out 120 | self.feature = feature 121 | self.ThreeDinit = ThreeDinit 122 | self.pooling = nn.AvgPool3d(kernel_size=3, padding=1, stride=1) 123 | self.business_layer.append(self.pooling) 124 | self.semantic_layer1 = nn.Sequential( 125 | Bottleneck3D(feature, feature // 4, bn_momentum=bn_momentum, expansion=4, stride=2, downsample= 126 | nn.Sequential( 127 | nn.AvgPool3d(kernel_size=2, stride=2), 128 | nn.Conv3d(feature, feature, 129 | kernel_size=1, stride=1, bias=False), 130 | norm_layer(feature, momentum=bn_momentum), 131 | ), norm_layer=norm_layer), 132 | Bottleneck3D(feature, feature // 4, bn_momentum=bn_momentum, norm_layer=norm_layer, dilation=[1, 1, 1]), 133 | Bottleneck3D(feature, feature // 4, bn_momentum=bn_momentum, norm_layer=norm_layer, dilation=[2, 2, 2]), 134 | Bottleneck3D(feature, feature // 4, bn_momentum=bn_momentum, norm_layer=norm_layer, dilation=[3, 3, 3]), 135 | ) 136 | self.business_layer.append(self.semantic_layer1) 137 | self.semantic_layer2 = nn.Sequential( 138 | Bottleneck3D(feature, feature // 4, bn_momentum=bn_momentum, expansion=8, stride=2, downsample= 139 | nn.Sequential( 140 | nn.AvgPool3d(kernel_size=2, stride=2), 141 | nn.Conv3d(feature, feature * 2, 142 | kernel_size=1, stride=1, bias=False), 143 | norm_layer(feature * 2, momentum=bn_momentum), 144 | ), norm_layer=norm_layer), 145 | Bottleneck3D(feature * 2, feature // 2, bn_momentum=bn_momentum, norm_layer=norm_layer, dilation=[1, 1, 1]), 146 | Bottleneck3D(feature * 2, feature // 2, bn_momentum=bn_momentum, norm_layer=norm_layer, dilation=[2, 2, 2]), 147 | Bottleneck3D(feature * 2, feature // 2, bn_momentum=bn_momentum, norm_layer=norm_layer, dilation=[3, 3, 3]), 148 | ) 149 | self.business_layer.append(self.semantic_layer2) 150 | self.classify_semantic = nn.ModuleList([ 151 | nn.Sequential( 152 | nn.ConvTranspose3d(feature * 2, feature, kernel_size=3, stride=2, padding=1, dilation=1, 153 | output_padding=1), 154 | norm_layer(feature, momentum=bn_momentum), 155 | nn.ReLU(inplace=False), 156 | ), 157 | nn.Sequential( 158 | nn.ConvTranspose3d(feature, feature, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1), 159 | norm_layer(feature, momentum=bn_momentum), 160 | nn.ReLU(inplace=False), 161 | ), 162 | nn.Sequential( 163 | nn.Dropout3d(.1), 164 | nn.Conv3d(feature, class_num, kernel_size=1, bias=True) 165 | ), 166 | nn.Sequential( 167 | nn.Dropout3d(.1), 168 | nn.Conv3d(feature, 2, kernel_size=1, bias=True) 169 | )] 170 | ) 171 | self.business_layer.append(self.classify_semantic) 172 | self.oper_tsdf = nn.Sequential( 173 | nn.Conv3d(1, 3, kernel_size=3, padding=1, bias=False), 174 | norm_layer(3, momentum=bn_momentum), 175 | nn.ReLU(), 176 | nn.Conv3d(3, 64, kernel_size=3, padding=1, bias=False), 177 | norm_layer(64, momentum=bn_momentum), 178 | nn.ReLU(), 179 | nn.Conv3d(64, feature, kernel_size=3, padding=1, bias=False), 180 | norm_layer(feature, momentum=bn_momentum), 181 | nn.ReLU(inplace=False), 182 | ) 183 | self.oper_raw = nn.Sequential( 184 | nn.Conv3d(12, 3, kernel_size=3, padding=1, bias=False), 185 | norm_layer(3, momentum=bn_momentum), 186 | nn.ReLU(), 187 | nn.Conv3d(3, 64, kernel_size=3, padding=1, bias=False), 188 | norm_layer(64, momentum=bn_momentum), 189 | nn.ReLU(), 190 | nn.Conv3d(64, feature, kernel_size=3, padding=1, bias=False), 191 | norm_layer(feature, momentum=bn_momentum), 192 | nn.ReLU(inplace=False), 193 | ) 194 | self.business_layer.append(self.oper_raw) 195 | self.business_layer.append(self.oper_tsdf) 196 | 197 | def forward(self, raw, tsdf): 198 | edge_raw = self.oper_raw(raw) 199 | edge_tsdf = self.oper_tsdf(tsdf) 200 | seg_fea = edge_tsdf + edge_raw 201 | semantic1 = self.semantic_layer1(seg_fea) + F.interpolate(seg_fea, size=[30, 18, 30]) 202 | semantic2 = self.semantic_layer2(semantic1) 203 | up_sem1 = self.classify_semantic[0](semantic2) 204 | up_sem1 = up_sem1 + semantic1 205 | up_sem2 = self.classify_semantic[1](up_sem1) 206 | up_sem2 = up_sem2 + F.interpolate(up_sem1, size=[60, 36, 60], mode="trilinear", align_corners=True) 207 | pred_semantic = self.classify_semantic[2](up_sem2) 208 | 209 | return pred_semantic, None 210 | 211 | class Network(nn.Module): 212 | def __init__(self, class_num, norm_layer, resnet_out=2048, feature=512, ThreeDinit=True, 213 | bn_momentum=0.1, pretrained_model=None, eval=False, freeze_bn=False): 214 | super(Network, self).__init__() 215 | self.business_layer = [] 216 | 217 | self.dilate = 2 218 | self.stage2 = STAGE2(class_num, norm_layer, resnet_out=resnet_out, feature=feature, ThreeDinit=ThreeDinit, 219 | bn_momentum=bn_momentum, pretrained_model=pretrained_model, eval=eval, freeze_bn=freeze_bn) 220 | self.business_layer += self.stage2.business_layer 221 | 222 | def forward(self, ttt, tsdf): 223 | pp, ppp = self.stage2(ttt, tsdf) 224 | 225 | if self.training: 226 | return pp, ppp 227 | return pp, ppp 228 | 229 | def _nostride_dilate(self, m, dilate): 230 | if isinstance(m, nn.Conv2d): 231 | if m.stride == (2, 2): 232 | m.stride = (1, 1) 233 | if m.kernel_size == (3, 3): 234 | m.dilation = (dilate, dilate) 235 | m.padding = (dilate, dilate) 236 | 237 | else: 238 | if m.kernel_size == (3, 3): 239 | m.dilation = (dilate, dilate) 240 | m.padding = (dilate, dilate) 241 | 242 | 243 | if __name__ == '__main__': 244 | pass 245 | -------------------------------------------------------------------------------- /exps/SemanticSceneCompletion/nyu.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2020/05 4 | # @Author : Xiaokang Chen 5 | 6 | import numpy as np 7 | import torch 8 | from datasets.BaseDataset import BaseDataset 9 | import os 10 | import cv2 11 | 12 | class NYUv2(BaseDataset): 13 | def __init__(self, setting, split_name, preprocess=None, 14 | file_length=None, s3client=None): 15 | super(NYUv2, self).__init__(setting, split_name, preprocess, file_length) 16 | self._split_name = split_name 17 | self._i_path = setting['i_root'] 18 | self._g_path = setting['g_root'] 19 | self._h_path = setting['h_root'] 20 | self._m_path = setting['m_root'] 21 | self._t_source = setting['t_source'] 22 | self._e_source = setting['e_source'] 23 | self._file_names = self._get_file_names(split_name) 24 | self._file_length = file_length 25 | self.preprocess = preprocess 26 | 27 | def _get_file_names(self, split_name): 28 | assert split_name in ['train', 'val'] 29 | source = self._t_source 30 | if split_name == "val": 31 | source = self._e_source 32 | 33 | file_names = [] 34 | with open(source) as f: 35 | files = f.readlines() 36 | 37 | for item in files: 38 | item = item.strip() 39 | item = item.split('\t') 40 | img_name = item[0] 41 | file_names.append([img_name, None]) 42 | 43 | return file_names 44 | 45 | 46 | def __getitem__(self, index): 47 | 48 | if self._file_length is not None: 49 | names = self._construct_new_file_names(self._file_length)[index] 50 | else: 51 | names = self._file_names[index] 52 | item_idx = names[0] 53 | t_path = os.path.join(self._i_path, 'voxel_bisenet_res34_4e-3/'+names[0]+'.npz') 54 | 55 | i_path = os.path.join(self._i_path, 'R', 'NYU'+item_idx+'_colors.png') 56 | g_path = os.path.join(self._g_path, 'L/'+item_idx+'.npz') 57 | m_path = os.path.join(self._m_path, item_idx+'.npz') 58 | w_path = os.path.join(self._i_path, 'T/'+item_idx+'.npz') 59 | item_name = item_idx 60 | ttt, www, mmm, ggg, tsdf = self._fetch_data(t_path, w_path, m_path, g_path) 61 | 62 | if self._split_name is 'train': 63 | ggg = torch.from_numpy(np.ascontiguousarray(ggg)).long() 64 | tsdf = torch.from_numpy(np.ascontiguousarray(tsdf)).float() 65 | mmm = torch.from_numpy(np.ascontiguousarray(mmm)).long() 66 | www = torch.from_numpy(np.ascontiguousarray(www)).float() 67 | ttt = torch.from_numpy(np.ascontiguousarray(ttt)).float() 68 | 69 | output_dict = dict(lll=ggg, tsdf=tsdf, www=www, mmm=mmm, 70 | ttt=ttt, fn=str(item_name), n=len(self._file_names)) 71 | return output_dict 72 | 73 | def _fetch_data(self, t_path, w_path, m_path, g_path,dtype=None): 74 | www = np.load(w_path)['arr_1'].astype(np.float32).reshape(60,36,60) 75 | _, file_extension = os.path.splitext(t_path) 76 | if file_extension in ['.npz']: 77 | ttt = np.load(t_path)['arr_0'].astype(np.int32) 78 | else: 79 | ttt = np.load(t_path).astype(np.int32) 80 | ttt = torch.LongTensor(ttt) 81 | ttt = self.get_one_hot( ttt, 12) 82 | ttt = ttt.permute(3,0,1,2) 83 | ggg = np.load(g_path)['arr_0'].astype(np.int64) 84 | tsdf = np.load(w_path)['arr_0'].astype(np.float32).reshape(1,60,36,60) 85 | mmm = np.load(m_path)['arr_0'].astype(np.int64) 86 | www = www.flatten() 87 | return ttt, www, mmm, ggg, tsdf 88 | 89 | def get_one_hot(self, label, N): ## N classes shapenet N=4, NYU N=8 90 | size = list(label.size()) 91 | label = label.view(-1) 92 | ones = torch.sparse.torch.eye(N) 93 | ones = ones.index_select(0, label) 94 | size.append(N) 95 | return ones.view(*size) 96 | @classmethod 97 | def get_class_colors(*args): 98 | def uint82bin(n, count=8): 99 | """returns the binary of integer n, count refers to amount of bits""" 100 | return ''.join([str((n >> y) & 1) for y in range(count - 1, -1, -1)]) 101 | 102 | N = 41 103 | cmap = np.zeros((N, 3), dtype=np.uint8) 104 | for i in range(N): 105 | r, g, b = 0, 0, 0 106 | id = i 107 | for j in range(7): 108 | str_id = uint82bin(id) 109 | r = r ^ (np.uint8(str_id[-1]) << (7 - j)) 110 | g = g ^ (np.uint8(str_id[-2]) << (7 - j)) 111 | b = b ^ (np.uint8(str_id[-3]) << (7 - j)) 112 | id = id >> 3 113 | cmap[i, 0] = r 114 | cmap[i, 1] = g 115 | cmap[i, 2] = b 116 | class_colors = cmap.tolist() 117 | return class_colors 118 | 119 | @classmethod 120 | def get_class_names(*args): 121 | 122 | return ['wall','floor','cabinet','bed','chair','sofa','table','door','window','bookshelf','picture','counter','blinds', 123 | 'desk','shelves','curtain','dresser','pillow','mirror','floor mat','clothes','ceiling','books','refridgerator', 124 | 'television','paper','towel','shower curtain','box','whiteboard','person','night stand','toilet', 125 | 'sink','lamp','bathtub','bag','otherstructure','otherfurniture','otherprop'] 126 | @classmethod 127 | def transform_label(cls, pred, name): 128 | label = np.zeros(pred.shape) 129 | ids = np.unique(pred) 130 | for id in ids: 131 | label[np.where(pred == id)] = cls.trans_labels[id] 132 | 133 | new_name = (name.split('.')[0]).split('_')[:-1] 134 | new_name = '_'.join(new_name) + '.png' 135 | 136 | return label, new_name 137 | -------------------------------------------------------------------------------- /exps/SemanticSceneCompletion/st-run.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | #export OMP_NUM_THREADS=4 3 | 4 | GPUNUM=16 5 | NODENUM=1 6 | JOBNAME=demo 7 | 8 | PART=vi_irdc 9 | 10 | TOOLS="srun --partition=$PART --gres=gpu:${GPUNUM} -n$NODENUM --ntasks-per-node=1" 11 | 12 | $TOOLS --job-name=$JOBNAME sh -c "python -m torch.distributed.launch --nnodes=$NODENUM --nproc_per_node=$GPUNUM --node_rank \$SLURM_PROCID --master_addr=\$(sinfo -Nh -n \$SLURM_NODELIST | head -n 1 | cut -d ' ' -f 1) --master_port 29058 train.py" 13 | 14 | -------------------------------------------------------------------------------- /exps/SemanticSceneCompletion/test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | #export OMP_NUM_THREADS=4 3 | 4 | GPUNUM=8 5 | NODENUM=1 6 | JOBNAME=test 7 | PART=vi_irdc 8 | 9 | TOOLS="srun --partition=$PART --gres=gpu:${GPUNUM} -n$NODENUM --ntasks-per-node=1" 10 | 11 | $TOOLS --job-name=$JOBNAME python eval.py -e 249-250 -d 0-1 12 | -------------------------------------------------------------------------------- /exps/SemanticSceneCompletion/train.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import os.path as osp 3 | import os 4 | import sys 5 | import time 6 | import argparse 7 | from tqdm import tqdm 8 | 9 | import torch 10 | import torch.nn as nn 11 | import torch.distributed as dist 12 | import torch.backends.cudnn as cudnn 13 | 14 | import numpy as np 15 | from config import config 16 | from dataloader import get_train_loader 17 | from network import Network 18 | from nyu import NYUv2 19 | from utils.init_func import init_weight, group_weight 20 | from engine.lr_policy import WarmUpPolyLR, PolyLR 21 | from engine.engine import Engine 22 | from seg_opr.sync_bn import DataParallelModel, Reduce, BatchNorm2d 23 | from tensorboardX import SummaryWriter 24 | 25 | try: 26 | from apex.parallel import DistributedDataParallel, SyncBatchNorm 27 | except ImportError: 28 | raise ImportError( 29 | "Please install apex from https://www.github.com/nvidia/apex .") 30 | 31 | parser = argparse.ArgumentParser() 32 | 33 | s3client = None 34 | 35 | port = str(int(float(time.time())) % 20) 36 | os.environ['MASTER_PORT'] = str(190802 + int(port)) 37 | 38 | with Engine(custom_parser=parser) as engine: 39 | args = parser.parse_args() 40 | cudnn.benchmark = True 41 | seed = config.seed 42 | if engine.distributed: 43 | seed = engine.local_rank 44 | torch.manual_seed(seed) 45 | if torch.cuda.is_available(): 46 | torch.cuda.manual_seed(seed) 47 | train_loader, train_sampler = get_train_loader(engine, NYUv2, s3client) 48 | if engine.distributed and (engine.local_rank == 0): 49 | tb_dir = config.tb_dir + '/{}'.format(time.strftime("%b%d_%d-%H-%M", time.localtime())) 50 | generate_tb_dir = config.tb_dir + '/tb' 51 | writer = SummaryWriter(log_dir=tb_dir) 52 | engine.link_tb(tb_dir, generate_tb_dir) 53 | criterion = nn.CrossEntropyLoss(reduction='mean', ignore_index=255) 54 | if engine.distributed: 55 | BatchNorm2d = SyncBatchNorm 56 | model = Network(class_num=config.num_classes, feature=128, bn_momentum=config.bn_momentum, 57 | norm_layer=BatchNorm2d) 58 | init_weight(model.business_layer, nn.init.kaiming_normal_, 59 | BatchNorm2d, config.bn_eps, config.bn_momentum, 60 | mode='fan_in') 61 | base_lr = config.lr 62 | if engine.distributed: 63 | base_lr = config.lr 64 | params_list = [] 65 | for module in model.business_layer: 66 | params_list = group_weight(params_list, module, BatchNorm2d,base_lr ) 67 | optimizer = torch.optim.SGD(params_list,lr=base_lr,momentum=config.momentum,weight_decay=config.weight_decay) 68 | total_iteration = config.nepochs * config.niters_per_epoch 69 | lr_policy = PolyLR(base_lr, config.lr_power, total_iteration) 70 | if engine.distributed: 71 | print('distributed !!') 72 | if torch.cuda.is_available(): 73 | model.cuda() 74 | model = DistributedDataParallel(model) 75 | else: 76 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 77 | model = DataParallelModel(model, device_ids=engine.devices) 78 | model.to(device) 79 | engine.register_state(dataloader=train_loader, model=model,optimizer=optimizer) 80 | if engine.continue_state_object: 81 | engine.restore_checkpoint() 82 | model.train() 83 | print('begin train') 84 | for epoch in range(engine.state.epoch, config.nepochs): 85 | if engine.distributed: 86 | train_sampler.set_epoch(epoch) 87 | bar_format = '{desc}[{elapsed}<{remaining},{rate_fmt}]' 88 | pbar = tqdm(range(config.niters_per_epoch), file=sys.stdout, 89 | bar_format=bar_format) 90 | dataloader = iter(train_loader) 91 | for idx in pbar: 92 | optimizer.zero_grad() 93 | engine.update_iteration(epoch, idx) 94 | minibatch = dataloader.next() 95 | lll = minibatch['lll'] 96 | tsdf= minibatch['tsdf'] 97 | www = minibatch['www'] 98 | ttt = minibatch['ttt'] 99 | mmm = minibatch['mmm'] 100 | 101 | lll = lll.cuda(non_blocking=True) 102 | ttt = ttt.cuda(non_blocking=True) 103 | www = www.cuda(non_blocking=True) 104 | mmm = mmm.cuda(non_blocking=True) 105 | tsdf = tsdf.cuda(non_blocking=True) 106 | 107 | output, boutput = model(ttt, tsdf) 108 | cri_weights = 10 * torch.FloatTensor([0.010820392313388523, 0.4585814244886793, 0.0411831291920445, 0.04826918042332931, 0.33162496143513115, 0.22373353821746247, 3*0.09748478737233816, 0.1478032329336482, 0.16258443136359715, 1.0, 0.0254366993244824, 0.05126348601814224]) 109 | criterion = nn.CrossEntropyLoss(ignore_index=255, reduction='none', 110 | weight=cri_weights).cuda() 111 | selectindex = torch.nonzero(www.view(-1)).view(-1) 112 | filterLabel = torch.index_select(lll.view(-1), 0, selectindex) 113 | filterOutput = torch.index_select(output.permute( 114 | 0, 2, 3, 4, 1).contiguous().view(-1, 12), 0, selectindex) 115 | loss_semantic = criterion(filterOutput, filterLabel) 116 | loss_semantic = torch.mean(loss_semantic) 117 | 118 | if engine.distributed: 119 | dist.all_reduce(loss_semantic, dist.ReduceOp.SUM) 120 | loss_semantic = loss_semantic / engine.world_size 121 | else: 122 | loss = Reduce.apply(*loss) / len(loss) 123 | 124 | current_idx = epoch * config.niters_per_epoch + idx 125 | lr = lr_policy.get_lr(current_idx) 126 | 127 | optimizer.param_groups[0]['lr'] = lr 128 | for i in range(1, len(optimizer.param_groups)): 129 | optimizer.param_groups[i]['lr'] = lr 130 | 131 | loss = loss_semantic 132 | loss.backward() 133 | nn.utils.clip_grad_norm_(model.parameters(), max_norm=20, norm_type=2) 134 | optimizer.step() 135 | print_str = 'Epoch{}/{}'.format(epoch, config.nepochs) \ 136 | + ' Iter{}/{}:'.format(idx + 1, config.niters_per_epoch) \ 137 | + ' lr=%.2e' % lr \ 138 | + ' loss=%.5f' % (loss.item()) 139 | 140 | pbar.set_description(print_str, refresh=False) 141 | if engine.distributed and (engine.local_rank == 0): 142 | writer.add_scalar("Loss/Epoch", loss.item(), epoch) 143 | 144 | if (epoch > config.nepochs // 4) and ((epoch % config.snapshot_iter == 0) or (epoch == config.nepochs - 1)): 145 | if engine.distributed and (engine.local_rank == 0): 146 | engine.save_and_link_checkpoint(config.snapshot_dir, 147 | config.log_dir, 148 | config.log_dir_link) 149 | elif not engine.distributed: 150 | engine.save_and_link_checkpoint(config.snapshot_dir, 151 | config.log_dir, 152 | config.log_dir_link) 153 | 154 | if engine.distributed and (engine.local_rank == 0): 155 | writer.close() 156 | -------------------------------------------------------------------------------- /figs/nyu_res.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/figs/nyu_res.png -------------------------------------------------------------------------------- /figs/nyu_vis.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/figs/nyu_vis.png -------------------------------------------------------------------------------- /furnace/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/.DS_Store -------------------------------------------------------------------------------- /furnace/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/__init__.py -------------------------------------------------------------------------------- /furnace/base_model/README.md: -------------------------------------------------------------------------------- 1 | # Base Model 2 | 3 | ## Performance 4 | 5 | | Model | Top-1 | Top-5 | Origin | 6 | |:------------:|:-----:|:-----:|:--------------------:| 7 | | resnet18_v1 | 70.93 | 89.92 | converted from gluon | 8 | | resnet34_v1 | 74.37 | 91.87 | converted from gluon | 9 | | resnet50_v1 | 76.47 | 93.13 | converted from gluon | 10 | | resnet101_v1 | 78.34 | 94.01 | converted from gluon | 11 | | resnet152_v1 | 79.00 | 94.38 | converted from gluon | 12 | | resnet18_v1c | | | converted from gluon | 13 | | resnet34_v1c | | | converted from gluon | 14 | | resnet50_v1c | | | converted from gluon | 15 | | resnet101_v1c| | | converted from gluon | 16 | | resnet152_v1c| | | converted from gluon | 17 | -------------------------------------------------------------------------------- /furnace/base_model/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnet import ResNet, resnet18, resnet34, resnet50, resnet101, resnet152, Bottleneck 2 | from .xception import Xception, xception39 -------------------------------------------------------------------------------- /furnace/base_model/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/base_model/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/base_model/__pycache__/resnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/base_model/__pycache__/resnet.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/base_model/resnet.py: -------------------------------------------------------------------------------- 1 | import functools 2 | import torch.nn as nn 3 | 4 | from utils.pyt_utils import load_model 5 | 6 | 7 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 8 | 'resnet152'] 9 | 10 | 11 | def conv3x3(in_planes, out_planes, stride=1): 12 | """3x3 convolution with padding""" 13 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 14 | padding=1, bias=False) 15 | 16 | 17 | class BasicBlock(nn.Module): 18 | expansion = 1 19 | 20 | def __init__(self, inplanes, planes, stride=1, norm_layer=None, 21 | bn_eps=1e-5, bn_momentum=0.1, downsample=None, inplace=True): 22 | super(BasicBlock, self).__init__() 23 | self.conv1 = conv3x3(inplanes, planes, stride) 24 | self.bn1 = norm_layer(planes, eps=bn_eps, momentum=bn_momentum) 25 | self.relu = nn.ReLU(inplace=inplace) 26 | self.relu_inplace = nn.ReLU(inplace=True) 27 | self.conv2 = conv3x3(planes, planes) 28 | self.bn2 = norm_layer(planes, eps=bn_eps, momentum=bn_momentum) 29 | self.downsample = downsample 30 | self.stride = stride 31 | self.inplace = inplace 32 | 33 | def forward(self, x): 34 | residual = x 35 | 36 | out = self.conv1(x) 37 | out = self.bn1(out) 38 | out = self.relu(out) 39 | 40 | out = self.conv2(out) 41 | out = self.bn2(out) 42 | 43 | if self.downsample is not None: 44 | residual = self.downsample(x) 45 | 46 | if self.inplace: 47 | out += residual 48 | else: 49 | out = out + residual 50 | 51 | out = self.relu_inplace(out) 52 | 53 | return out 54 | 55 | 56 | class Bottleneck(nn.Module): 57 | expansion = 4 58 | 59 | def __init__(self, inplanes, planes, stride=1, 60 | norm_layer=None, bn_eps=1e-5, bn_momentum=0.1, 61 | downsample=None, inplace=True): 62 | super(Bottleneck, self).__init__() 63 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 64 | self.bn1 = norm_layer(planes, eps=bn_eps, momentum=bn_momentum) 65 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 66 | padding=1, bias=False) 67 | self.bn2 = norm_layer(planes, eps=bn_eps, momentum=bn_momentum) 68 | self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, 69 | bias=False) 70 | self.bn3 = norm_layer(planes * self.expansion, eps=bn_eps, 71 | momentum=bn_momentum) 72 | self.relu = nn.ReLU(inplace=inplace) 73 | self.relu_inplace = nn.ReLU(inplace=True) 74 | self.downsample = downsample 75 | self.stride = stride 76 | self.inplace = inplace 77 | 78 | def forward(self, x): 79 | residual = x 80 | 81 | out = self.conv1(x) 82 | out = self.bn1(out) 83 | out = self.relu(out) 84 | 85 | out = self.conv2(out) 86 | out = self.bn2(out) 87 | out = self.relu(out) 88 | 89 | out = self.conv3(out) 90 | out = self.bn3(out) 91 | 92 | if self.downsample is not None: 93 | residual = self.downsample(x) 94 | 95 | if self.inplace: 96 | out += residual 97 | else: 98 | out = out + residual 99 | out = self.relu_inplace(out) 100 | 101 | return out 102 | 103 | 104 | class ResNet(nn.Module): 105 | 106 | def __init__(self, block, layers, norm_layer=nn.BatchNorm2d, bn_eps=1e-5, 107 | bn_momentum=0.1, deep_stem=False, stem_width=32, inplace=True): 108 | self.inplanes = stem_width * 2 if deep_stem else 64 109 | super(ResNet, self).__init__() 110 | if deep_stem: 111 | self.conv1 = nn.Sequential( 112 | nn.Conv2d(3, stem_width, kernel_size=3, stride=2, padding=1, 113 | bias=False), 114 | norm_layer(stem_width, eps=bn_eps, momentum=bn_momentum), 115 | nn.ReLU(inplace=inplace), 116 | nn.Conv2d(stem_width, stem_width, kernel_size=3, stride=1, 117 | padding=1, 118 | bias=False), 119 | norm_layer(stem_width, eps=bn_eps, momentum=bn_momentum), 120 | nn.ReLU(inplace=inplace), 121 | nn.Conv2d(stem_width, stem_width * 2, kernel_size=3, stride=1, 122 | padding=1, 123 | bias=False), 124 | ) 125 | else: 126 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 127 | bias=False) 128 | 129 | self.bn1 = norm_layer(stem_width * 2 if deep_stem else 64, eps=bn_eps, 130 | momentum=bn_momentum) 131 | self.relu = nn.ReLU(inplace=inplace) 132 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 133 | self.layer1 = self._make_layer(block, norm_layer, 64, layers[0], 134 | inplace, 135 | bn_eps=bn_eps, bn_momentum=bn_momentum) 136 | self.layer2 = self._make_layer(block, norm_layer, 128, layers[1], 137 | inplace, stride=2, 138 | bn_eps=bn_eps, bn_momentum=bn_momentum) 139 | self.layer3 = self._make_layer(block, norm_layer, 256, layers[2], 140 | inplace, stride=2, 141 | bn_eps=bn_eps, bn_momentum=bn_momentum) 142 | self.layer4 = self._make_layer(block, norm_layer, 512, layers[3], 143 | inplace, stride=2, 144 | bn_eps=bn_eps, bn_momentum=bn_momentum) 145 | 146 | def _make_layer(self, block, norm_layer, planes, blocks, inplace=True, 147 | stride=1, bn_eps=1e-5, bn_momentum=0.1): 148 | downsample = None 149 | if stride != 1 or self.inplanes != planes * block.expansion: 150 | downsample = nn.Sequential( 151 | nn.Conv2d(self.inplanes, planes * block.expansion, 152 | kernel_size=1, stride=stride, bias=False), 153 | norm_layer(planes * block.expansion, eps=bn_eps, 154 | momentum=bn_momentum), 155 | ) 156 | 157 | layers = [] 158 | layers.append(block(self.inplanes, planes, stride, norm_layer, bn_eps, 159 | bn_momentum, downsample, inplace)) 160 | self.inplanes = planes * block.expansion 161 | for i in range(1, blocks): 162 | layers.append(block(self.inplanes, planes, 163 | norm_layer=norm_layer, bn_eps=bn_eps, 164 | bn_momentum=bn_momentum, inplace=inplace)) 165 | 166 | return nn.Sequential(*layers) 167 | 168 | def forward(self, x): 169 | x = self.conv1(x) 170 | x = self.bn1(x) 171 | x = self.relu(x) 172 | x = self.maxpool(x) 173 | 174 | blocks = [] 175 | x = self.layer1(x); 176 | blocks.append(x) 177 | x = self.layer2(x); 178 | blocks.append(x) 179 | x = self.layer3(x); 180 | blocks.append(x) 181 | x = self.layer4(x); 182 | blocks.append(x) 183 | 184 | return blocks 185 | 186 | 187 | def resnet18(pretrained_model=None, **kwargs): 188 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 189 | 190 | if pretrained_model is not None: 191 | model = load_model(model, pretrained_model) 192 | return model 193 | 194 | 195 | def resnet34(pretrained_model=None, **kwargs): 196 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 197 | 198 | if pretrained_model is not None: 199 | model = load_model(model, pretrained_model) 200 | return model 201 | 202 | 203 | def resnet50(pretrained_model=None, **kwargs): 204 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 205 | 206 | if pretrained_model is not None: 207 | model = load_model(model, pretrained_model) 208 | return model 209 | 210 | 211 | def resnet101(pretrained_model=None, **kwargs): 212 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 213 | 214 | if pretrained_model is not None: 215 | model = load_model(model, pretrained_model) 216 | return model 217 | 218 | 219 | def resnet152(pretrained_model=None, **kwargs): 220 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 221 | 222 | if pretrained_model is not None: 223 | model = load_model(model, pretrained_model) 224 | return model 225 | -------------------------------------------------------------------------------- /furnace/base_model/xception.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, division, absolute_import 2 | import torch.nn as nn 3 | 4 | from seg_opr.seg_oprs import ConvBnRelu 5 | from utils.pyt_utils import load_model 6 | 7 | __all__ = ['Xception', 'xception39'] 8 | 9 | 10 | class SeparableConvBnRelu(nn.Module): 11 | def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, 12 | padding=0, dilation=1, 13 | has_relu=True, norm_layer=nn.BatchNorm2d): 14 | super(SeparableConvBnRelu, self).__init__() 15 | 16 | self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, stride, 17 | padding, dilation, groups=in_channels, 18 | bias=False) 19 | self.point_wise_cbr = ConvBnRelu(in_channels, out_channels, 1, 1, 0, 20 | has_bn=True, norm_layer=norm_layer, 21 | has_relu=has_relu, has_bias=False) 22 | 23 | def forward(self, x): 24 | x = self.conv1(x) 25 | x = self.point_wise_cbr(x) 26 | return x 27 | 28 | 29 | class Block(nn.Module): 30 | expansion = 4 31 | 32 | def __init__(self, in_channels, mid_out_channels, has_proj, stride, 33 | dilation=1, norm_layer=nn.BatchNorm2d): 34 | super(Block, self).__init__() 35 | self.has_proj = has_proj 36 | 37 | if has_proj: 38 | self.proj = SeparableConvBnRelu(in_channels, 39 | mid_out_channels * self.expansion, 40 | 3, stride, 1, 41 | has_relu=False, 42 | norm_layer=norm_layer) 43 | 44 | self.residual_branch = nn.Sequential( 45 | SeparableConvBnRelu(in_channels, mid_out_channels, 46 | 3, stride, dilation, dilation, 47 | has_relu=True, norm_layer=norm_layer), 48 | SeparableConvBnRelu(mid_out_channels, mid_out_channels, 3, 1, 1, 49 | has_relu=True, norm_layer=norm_layer), 50 | SeparableConvBnRelu(mid_out_channels, 51 | mid_out_channels * self.expansion, 3, 1, 1, 52 | has_relu=False, norm_layer=norm_layer)) 53 | self.relu = nn.ReLU(inplace=True) 54 | 55 | def forward(self, x): 56 | shortcut = x 57 | if self.has_proj: 58 | shortcut = self.proj(x) 59 | 60 | residual = self.residual_branch(x) 61 | output = self.relu(shortcut + residual) 62 | 63 | return output 64 | 65 | 66 | class Xception(nn.Module): 67 | def __init__(self, block, layers, channels, norm_layer=nn.BatchNorm2d): 68 | super(Xception, self).__init__() 69 | 70 | self.in_channels = 8 71 | self.conv1 = ConvBnRelu(3, self.in_channels, 3, 2, 1, 72 | has_bn=True, norm_layer=norm_layer, 73 | has_relu=True, has_bias=False) 74 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 75 | 76 | self.layer1 = self._make_layer(block, norm_layer, 77 | layers[0], channels[0], stride=2) 78 | self.layer2 = self._make_layer(block, norm_layer, 79 | layers[1], channels[1], stride=2) 80 | self.layer3 = self._make_layer(block, norm_layer, 81 | layers[2], channels[2], stride=2) 82 | 83 | def _make_layer(self, block, norm_layer, blocks, 84 | mid_out_channels, stride=1): 85 | layers = [] 86 | has_proj = True if stride > 1 else False 87 | layers.append(block(self.in_channels, mid_out_channels, has_proj, 88 | stride=stride, norm_layer=norm_layer)) 89 | self.in_channels = mid_out_channels * block.expansion 90 | for i in range(1, blocks): 91 | layers.append(block(self.in_channels, mid_out_channels, 92 | has_proj=False, stride=1, 93 | norm_layer=norm_layer)) 94 | 95 | return nn.Sequential(*layers) 96 | 97 | def forward(self, x): 98 | x = self.conv1(x) 99 | x = self.maxpool(x) 100 | 101 | blocks = [] 102 | x = self.layer1(x); 103 | blocks.append(x) 104 | x = self.layer2(x); 105 | blocks.append(x) 106 | x = self.layer3(x); 107 | blocks.append(x) 108 | 109 | return blocks 110 | 111 | 112 | def xception39(pretrained_model=None, **kwargs): 113 | model = Xception(Block, [4, 8, 4], [16, 32, 64], **kwargs) 114 | 115 | if pretrained_model is not None: 116 | model = load_model(model, pretrained_model) 117 | return model 118 | -------------------------------------------------------------------------------- /furnace/datasets/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/datasets/.DS_Store -------------------------------------------------------------------------------- /furnace/datasets/BaseDataset.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2017/12/16 下午8:41 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : BaseDataset.py 7 | 8 | import os 9 | import time 10 | import cv2 11 | import torch 12 | import numpy as np 13 | 14 | import torch.utils.data as data 15 | 16 | 17 | class BaseDataset(data.Dataset): 18 | def __init__(self, setting, split_name, preprocess=None, 19 | file_length=None): 20 | super(BaseDataset, self).__init__() 21 | 22 | def __len__(self): 23 | if self._file_length is not None: 24 | return self._file_length 25 | return len(self._file_names) 26 | 27 | def _get_file_names(self, split_name, train_extra=False): 28 | assert split_name in ['train', 'val'] 29 | source = self._train_source 30 | if split_name == "val": 31 | source = self._eval_source 32 | 33 | file_names = [] 34 | with open(source) as f: 35 | files = f.readlines() 36 | 37 | for item in files: 38 | img_name, gt_name = self._process_item_names(item) 39 | file_names.append([img_name, gt_name]) 40 | 41 | if train_extra: 42 | file_names2 = [] 43 | source2 = self._train_source.replace('train', 'train_extra') 44 | with open(source2) as f: 45 | files2 = f.readlines() 46 | 47 | for item in files2: 48 | img_name, gt_name = self._process_item_names(item) 49 | file_names2.append([img_name, gt_name]) 50 | 51 | return file_names, file_names2 52 | 53 | return file_names 54 | 55 | def _construct_new_file_names(self, length): 56 | assert isinstance(length, int) 57 | files_len = len(self._file_names) # 原来一轮迭代的长度 58 | new_file_names = self._file_names * (length // files_len) # 按照设定获得的一轮迭代的长度 59 | 60 | rand_indices = torch.randperm(files_len).tolist() 61 | new_indices = rand_indices[:length % files_len] 62 | 63 | new_file_names += [self._file_names[i] for i in new_indices] 64 | 65 | return new_file_names 66 | 67 | @staticmethod 68 | def _process_item_names(item): 69 | item = item.strip() 70 | item = item.split('\t') 71 | img_name = item[0] 72 | gt_name = item[1] 73 | 74 | return img_name, gt_name 75 | 76 | def get_length(self): 77 | return self.__len__() 78 | 79 | @staticmethod 80 | def _open_image(filepath, mode=cv2.IMREAD_COLOR, dtype=None): 81 | # cv2: B G R 82 | # h w c 83 | img = np.array(cv2.imread(filepath, mode), dtype=dtype) 84 | 85 | return img 86 | 87 | @classmethod 88 | def get_class_colors(*args): 89 | raise NotImplementedError 90 | 91 | @classmethod 92 | def get_class_names(*args): 93 | raise NotImplementedError 94 | 95 | 96 | if __name__ == "__main__": 97 | data_setting = {'img_root': '', 98 | 'gt_root': '', 99 | 'train_source': '', 100 | 'eval_source': ''} 101 | bd = BaseDataset(data_setting, 'train', None) 102 | print(bd.get_class_names()) 103 | -------------------------------------------------------------------------------- /furnace/datasets/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/datasets/__init__.py -------------------------------------------------------------------------------- /furnace/datasets/__pycache__/BaseDataset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/datasets/__pycache__/BaseDataset.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/datasets/__pycache__/BaseDataset.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/datasets/__pycache__/BaseDataset.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/datasets/__pycache__/BaseDataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/datasets/__pycache__/BaseDataset.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/datasets/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/datasets/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/datasets/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/datasets/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/datasets/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/datasets/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/engine/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__init__.py -------------------------------------------------------------------------------- /furnace/engine/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/engine.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/engine.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/engine.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/engine.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/engine.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/engine.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/evaluator.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/evaluator.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/evaluator.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/evaluator.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/logger.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/logger.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/logger.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/logger.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/logger.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/logger.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/lr_policy.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/lr_policy.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/lr_policy.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/lr_policy.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/lr_policy.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/lr_policy.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/version.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/version.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/version.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/version.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/engine/__pycache__/version.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/engine/__pycache__/version.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/engine/engine.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/8/2 下午3:23 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : engine.py 7 | import os 8 | import os.path as osp 9 | import time 10 | import argparse 11 | 12 | import torch 13 | import torch.distributed as dist 14 | 15 | from .logger import get_logger 16 | from .version import __version__ 17 | from utils.pyt_utils import load_model, parse_devices, extant_file, link_file, \ 18 | ensure_dir 19 | 20 | logger = get_logger() 21 | 22 | 23 | class State(object): 24 | def __init__(self): 25 | self.epoch = 0 26 | self.iteration = 0 27 | self.dataloader = None 28 | self.model = None 29 | self.optimizer = None 30 | 31 | def register(self, **kwargs): 32 | for k, v in kwargs.items(): 33 | assert k in ['epoch', 'iteration', 'dataloader', 'model', 34 | 'optimizer'] 35 | setattr(self, k, v) 36 | 37 | 38 | class Engine(object): 39 | def __init__(self, custom_parser=None): 40 | self.version = __version__ 41 | logger.info( 42 | "PyTorch Version {}, Furnace Version {}".format(torch.__version__, 43 | self.version)) 44 | self.state = State() 45 | self.devices = None 46 | self.distributed = False 47 | 48 | if custom_parser is None: 49 | self.parser = argparse.ArgumentParser() 50 | else: 51 | assert isinstance(custom_parser, argparse.ArgumentParser) 52 | self.parser = custom_parser 53 | 54 | self.inject_default_parser() 55 | self.args = self.parser.parse_args() 56 | 57 | self.continue_state_object = self.args.continue_fpath 58 | 59 | if 'WORLD_SIZE' in os.environ: 60 | self.distributed = int(os.environ['WORLD_SIZE']) >= 1 61 | 62 | if self.distributed: 63 | self.local_rank = self.args.local_rank 64 | self.world_size = int(os.environ['WORLD_SIZE']) 65 | torch.cuda.set_device(self.local_rank) 66 | os.environ['MASTER_PORT'] = self.args.port 67 | dist.init_process_group(backend="nccl", init_method='env://') 68 | self.devices = [i for i in range(self.world_size)] 69 | else: 70 | self.devices = parse_devices(self.args.devices) 71 | 72 | def inject_default_parser(self): 73 | p = self.parser 74 | p.add_argument('-d', '--devices', default='', 75 | help='set data parallel training') 76 | p.add_argument('-c', '--continue', type=extant_file, 77 | metavar="FILE", 78 | dest="continue_fpath", 79 | help='continue from one certain checkpoint') 80 | p.add_argument('--local_rank', default=0, type=int, 81 | help='process rank on node') 82 | p.add_argument('-p', '--port', type=str, 83 | default='16001', 84 | dest="port", 85 | help='port for init_process_group') 86 | 87 | def register_state(self, **kwargs): 88 | self.state.register(**kwargs) 89 | 90 | def update_iteration(self, epoch, iteration): 91 | self.state.epoch = epoch 92 | self.state.iteration = iteration 93 | 94 | def save_checkpoint(self, path): 95 | logger.info("Saving checkpoint to file {}".format(path)) 96 | t_start = time.time() 97 | 98 | state_dict = {} 99 | 100 | from collections import OrderedDict 101 | new_state_dict = OrderedDict() 102 | for k, v in self.state.model.state_dict().items(): 103 | key = k 104 | if k.split('.')[0] == 'module': 105 | key = k[7:] 106 | new_state_dict[key] = v 107 | state_dict['model'] = new_state_dict 108 | state_dict['optimizer'] = self.state.optimizer.state_dict() 109 | state_dict['epoch'] = self.state.epoch 110 | state_dict['iteration'] = self.state.iteration 111 | 112 | t_iobegin = time.time() 113 | torch.save(state_dict, path) 114 | del state_dict 115 | del new_state_dict 116 | t_end = time.time() 117 | logger.info( 118 | "Save checkpoint to file {}, " 119 | "Time usage:\n\tprepare snapshot: {}, IO: {}".format( 120 | path, t_iobegin - t_start, t_end - t_iobegin)) 121 | 122 | def link_tb(self, source, target): 123 | ensure_dir(source) 124 | ensure_dir(target) 125 | link_file(source, target) 126 | 127 | 128 | def save_and_link_checkpoint(self, snapshot_dir, log_dir, log_dir_link): 129 | ensure_dir(snapshot_dir) 130 | if not osp.exists(log_dir_link): 131 | link_file(log_dir, log_dir_link) 132 | current_epoch_checkpoint = osp.join(snapshot_dir, 'epoch-{}.pth'.format( 133 | self.state.epoch)) 134 | self.save_checkpoint(current_epoch_checkpoint) 135 | last_epoch_checkpoint = osp.join(snapshot_dir, 136 | 'epoch-last.pth') 137 | link_file(current_epoch_checkpoint, last_epoch_checkpoint) 138 | 139 | def restore_checkpoint(self): 140 | t_start = time.time() 141 | if self.distributed: 142 | tmp = torch.load(self.continue_state_object, 143 | map_location=lambda storage, loc: storage.cuda( 144 | self.local_rank)) 145 | else: 146 | tmp = torch.load(self.continue_state_object) 147 | t_ioend = time.time() 148 | 149 | self.state.model = load_model(self.state.model, tmp['model'], 150 | True) 151 | self.state.optimizer.load_state_dict(tmp['optimizer']) 152 | self.state.epoch = tmp['epoch'] + 1 153 | self.state.iteration = tmp['iteration'] 154 | del tmp 155 | t_end = time.time() 156 | logger.info( 157 | "Load checkpoint from file {}, " 158 | "Time usage:\n\tIO: {}, restore snapshot: {}".format( 159 | self.continue_state_object, t_ioend - t_start, t_end - t_ioend)) 160 | 161 | def __enter__(self): 162 | return self 163 | 164 | def __exit__(self, type, value, tb): 165 | torch.cuda.empty_cache() 166 | if type is not None: 167 | logger.warning( 168 | "A exception occurred during Engine initialization, " 169 | "give up running process") 170 | return False 171 | -------------------------------------------------------------------------------- /furnace/engine/logger.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/8/2 上午11:48 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : logger.py 7 | import os 8 | import sys 9 | import logging 10 | 11 | from utils import pyt_utils 12 | # from utils.pyt_utils import ensure_dir 13 | 14 | _default_level_name = os.getenv('ENGINE_LOGGING_LEVEL', 'INFO') 15 | _default_level = logging.getLevelName(_default_level_name.upper()) 16 | 17 | 18 | class LogFormatter(logging.Formatter): 19 | log_fout = None 20 | date_full = '[%(asctime)s %(lineno)d@%(filename)s:%(name)s] ' 21 | date = '%(asctime)s ' 22 | msg = '%(message)s' 23 | 24 | def format(self, record): 25 | if record.levelno == logging.DEBUG: 26 | mcl, mtxt = self._color_dbg, 'DBG' 27 | elif record.levelno == logging.WARNING: 28 | mcl, mtxt = self._color_warn, 'WRN' 29 | elif record.levelno == logging.ERROR: 30 | mcl, mtxt = self._color_err, 'ERR' 31 | else: 32 | mcl, mtxt = self._color_normal, '' 33 | 34 | if mtxt: 35 | mtxt += ' ' 36 | 37 | if self.log_fout: 38 | self.__set_fmt(self.date_full + mtxt + self.msg) 39 | formatted = super(LogFormatter, self).format(record) 40 | # self.log_fout.write(formatted) 41 | # self.log_fout.write('\n') 42 | # self.log_fout.flush() 43 | return formatted 44 | 45 | self.__set_fmt(self._color_date(self.date) + mcl(mtxt + self.msg)) 46 | formatted = super(LogFormatter, self).format(record) 47 | 48 | return formatted 49 | 50 | if sys.version_info.major < 3: 51 | def __set_fmt(self, fmt): 52 | self._fmt = fmt 53 | else: 54 | def __set_fmt(self, fmt): 55 | self._style._fmt = fmt 56 | 57 | @staticmethod 58 | def _color_dbg(msg): 59 | return '\x1b[36m{}\x1b[0m'.format(msg) 60 | 61 | @staticmethod 62 | def _color_warn(msg): 63 | return '\x1b[1;31m{}\x1b[0m'.format(msg) 64 | 65 | @staticmethod 66 | def _color_err(msg): 67 | return '\x1b[1;4;31m{}\x1b[0m'.format(msg) 68 | 69 | @staticmethod 70 | def _color_omitted(msg): 71 | return '\x1b[35m{}\x1b[0m'.format(msg) 72 | 73 | @staticmethod 74 | def _color_normal(msg): 75 | return msg 76 | 77 | @staticmethod 78 | def _color_date(msg): 79 | return '\x1b[32m{}\x1b[0m'.format(msg) 80 | 81 | 82 | def get_logger(log_dir=None, log_file=None, formatter=LogFormatter): 83 | logger = logging.getLogger() 84 | logger.setLevel(_default_level) 85 | del logger.handlers[:] 86 | 87 | if log_dir and log_file: 88 | pyt_utils.ensure_dir(log_dir) 89 | LogFormatter.log_fout = True 90 | file_handler = logging.FileHandler(log_file, mode='a') 91 | file_handler.setLevel(logging.INFO) 92 | file_handler.setFormatter(formatter) 93 | logger.addHandler(file_handler) 94 | 95 | stream_handler = logging.StreamHandler() 96 | stream_handler.setFormatter(formatter(datefmt='%d %H:%M:%S')) 97 | stream_handler.setLevel(0) 98 | logger.addHandler(stream_handler) 99 | return logger 100 | -------------------------------------------------------------------------------- /furnace/engine/lr_policy.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/8/1 上午1:50 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : lr_policy.py.py 7 | 8 | from abc import ABCMeta, abstractmethod 9 | 10 | 11 | class BaseLR(): 12 | __metaclass__ = ABCMeta 13 | 14 | @abstractmethod 15 | def get_lr(self, cur_iter): pass 16 | 17 | 18 | class PolyLR(BaseLR): 19 | def __init__(self, start_lr, lr_power, total_iters): 20 | self.start_lr = start_lr 21 | self.lr_power = lr_power 22 | self.total_iters = total_iters + 0.0 23 | 24 | def get_lr(self, cur_iter): 25 | return self.start_lr * ( 26 | (1 - float(cur_iter) / self.total_iters) ** self.lr_power) 27 | 28 | class WarmUpPolyLR(BaseLR): 29 | def __init__(self, start_lr, lr_power, total_iters, warmup_steps): 30 | self.start_lr = start_lr 31 | self.lr_power = lr_power 32 | self.total_iters = total_iters + 0.0 33 | self.warmup_steps = warmup_steps 34 | 35 | def get_lr(self, cur_iter): 36 | if cur_iter < self.warmup_steps: 37 | return self.start_lr * (cur_iter / self.warmup_steps) 38 | else: 39 | return self.start_lr * ( 40 | (1 - float(cur_iter) / self.total_iters) ** self.lr_power) 41 | 42 | class MultiStageLR(BaseLR): 43 | def __init__(self, lr_stages): 44 | assert type(lr_stages) in [list, tuple] and len(lr_stages[0]) == 2, \ 45 | 'lr_stages must be list or tuple, with [iters, lr] format' 46 | self._lr_stagess = lr_stages 47 | 48 | def get_lr(self, epoch): 49 | for it_lr in self._lr_stagess: 50 | if epoch < it_lr[0]: 51 | return it_lr[1] 52 | 53 | 54 | class LinearIncreaseLR(BaseLR): 55 | def __init__(self, start_lr, end_lr, warm_iters): 56 | self._start_lr = start_lr 57 | self._end_lr = end_lr 58 | self._warm_iters = warm_iters 59 | self._delta_lr = (end_lr - start_lr) / warm_iters 60 | 61 | def get_lr(self, cur_epoch): 62 | return self._start_lr + cur_epoch * self._delta_lr 63 | 64 | 65 | -------------------------------------------------------------------------------- /furnace/engine/version.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/8/3 下午2:59 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : version.py 7 | 8 | __version__ = '0.1.1' -------------------------------------------------------------------------------- /furnace/seg_opr/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/.DS_Store -------------------------------------------------------------------------------- /furnace/seg_opr/__init__.py: -------------------------------------------------------------------------------- 1 | from .seg_oprs import * 2 | -------------------------------------------------------------------------------- /furnace/seg_opr/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/__pycache__/metric.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/__pycache__/metric.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/__pycache__/metric.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/__pycache__/metric.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/__pycache__/seg_oprs.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/__pycache__/seg_oprs.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/__pycache__/seg_oprs.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/__pycache__/seg_oprs.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/__pycache__/seg_oprs.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/__pycache__/seg_oprs.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/metric.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | 3 | import numpy as np 4 | 5 | np.seterr(divide='ignore', invalid='ignore') 6 | 7 | 8 | # voc cityscapes metric 9 | def hist_info(n_cl, pred, gt): 10 | assert (pred.shape == gt.shape) 11 | k = (gt >= 0) & (gt < n_cl) 12 | labeled = np.sum(k) 13 | correct = np.sum((pred[k] == gt[k])) 14 | 15 | return np.bincount(n_cl * gt[k].astype(int) + pred[k].astype(int), 16 | minlength=n_cl ** 2).reshape(n_cl, 17 | n_cl), labeled, correct 18 | 19 | 20 | def compute_score(hist, correct, labeled): 21 | iu = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist)) 22 | mean_IU = np.nanmean(iu) 23 | mean_IU_no_back = np.nanmean(iu[1:]) 24 | freq = hist.sum(1) / hist.sum() 25 | freq_IU = (iu[freq > 0] * freq[freq > 0]).sum() 26 | mean_pixel_acc = correct / labeled 27 | 28 | return iu, mean_IU, mean_IU_no_back, mean_pixel_acc 29 | 30 | 31 | # ade metric 32 | def meanIoU(area_intersection, area_union): 33 | iou = 1.0 * np.sum(area_intersection, axis=1) / np.sum(area_union, axis=1) 34 | meaniou = np.nanmean(iou) 35 | meaniou_no_back = np.nanmean(iou[1:]) 36 | 37 | return iou, meaniou, meaniou_no_back 38 | 39 | 40 | def intersectionAndUnion(imPred, imLab, numClass): 41 | # Remove classes from unlabeled pixels in gt image. 42 | # We should not penalize detections in unlabeled portions of the image. 43 | imPred = imPred * (imLab >= 0) 44 | 45 | # Compute area intersection: 46 | intersection = imPred * (imPred == imLab) 47 | (area_intersection, _) = np.histogram(intersection, bins=numClass, 48 | range=(1, numClass)) 49 | 50 | # Compute area union: 51 | (area_pred, _) = np.histogram(imPred, bins=numClass, range=(1, numClass)) 52 | (area_lab, _) = np.histogram(imLab, bins=numClass, range=(1, numClass)) 53 | area_union = area_pred + area_lab - area_intersection 54 | 55 | return area_intersection, area_union 56 | 57 | 58 | def mean_pixel_accuracy(pixel_correct, pixel_labeled): 59 | mean_pixel_accuracy = 1.0 * np.sum(pixel_correct) / ( 60 | np.spacing(1) + np.sum(pixel_labeled)) 61 | 62 | return mean_pixel_accuracy 63 | 64 | 65 | def pixelAccuracy(imPred, imLab): 66 | # Remove classes from unlabeled pixels in gt image. 67 | # We should not penalize detections in unlabeled portions of the image. 68 | pixel_labeled = np.sum(imLab >= 0) 69 | pixel_correct = np.sum((imPred == imLab) * (imLab >= 0)) 70 | pixel_accuracy = 1.0 * pixel_correct / pixel_labeled 71 | 72 | return pixel_accuracy, pixel_correct, pixel_labeled 73 | -------------------------------------------------------------------------------- /furnace/seg_opr/parallel/parallel_apply.py: -------------------------------------------------------------------------------- 1 | # import threading 2 | import torch 3 | import torch.multiprocessing as mp 4 | from torch.cuda._utils import _get_device_index 5 | 6 | 7 | def get_a_var(obj): 8 | if isinstance(obj, torch.Tensor): 9 | return obj 10 | 11 | if isinstance(obj, list) or isinstance(obj, tuple): 12 | for result in map(get_a_var, obj): 13 | if isinstance(result, torch.Tensor): 14 | return result 15 | if isinstance(obj, dict): 16 | for result in map(get_a_var, obj.items()): 17 | if isinstance(result, torch.Tensor): 18 | return result 19 | return None 20 | 21 | 22 | def parallel_apply(modules, inputs, kwargs_tup=None, devices=None): 23 | r"""Applies each `module` in :attr:`modules` in parallel on arguments 24 | contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword) 25 | on each of :attr:`devices`. 26 | Args: 27 | modules (Module): modules to be parallelized 28 | inputs (tensor): inputs to the modules 29 | devices (list of int or torch.device): CUDA devices 30 | :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and 31 | :attr:`devices` (if given) should all have same length. Moreover, each 32 | element of :attr:`inputs` can either be a single object as the only argument 33 | to a module, or a collection of positional arguments. 34 | """ 35 | assert len(modules) == len(inputs) 36 | if kwargs_tup is not None: 37 | assert len(modules) == len(kwargs_tup) 38 | else: 39 | kwargs_tup = ({},) * len(modules) 40 | if devices is not None: 41 | assert len(modules) == len(devices) 42 | else: 43 | devices = [None] * len(modules) 44 | devices = list(map(lambda x: _get_device_index(x, True), devices)) 45 | context = mp.get_context('spawn') 46 | # lock = threading.Lock() 47 | # results = {} 48 | # results = [] 49 | results_queue = context.Queue(len(devices)) 50 | grad_enabled = torch.is_grad_enabled() 51 | 52 | def _worker(i, module, input, kwargs, device=None): 53 | torch.set_grad_enabled(grad_enabled) 54 | if device is None: 55 | device = get_a_var(input).get_device() 56 | try: 57 | with torch.cuda.device(device): 58 | # this also avoids accidental slicing of `input` if it is a Tensor 59 | if not isinstance(input, (list, tuple)): 60 | input = (input,) 61 | output = module(*input, **kwargs) 62 | results_queue.put(output) 63 | # with lock: 64 | # results[i] = output 65 | except Exception as e: 66 | results_queue.put(e) 67 | # with lock: 68 | # results[i] = e 69 | 70 | if len(modules) > 1: 71 | processes = [context.Process(target=_worker, 72 | args=(i, module, input, kwargs, device)) 73 | for i, (module, input, kwargs, device) in 74 | enumerate(zip(modules, inputs, kwargs_tup, devices))] 75 | 76 | for process in processes: 77 | process.start() 78 | for process in processes: 79 | process.join() 80 | else: 81 | _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0]) 82 | 83 | outputs = [] 84 | for i in range(len(inputs)): 85 | output = results_queue.get() 86 | if isinstance(output, Exception): 87 | raise output 88 | outputs.append(output) 89 | return outputs 90 | -------------------------------------------------------------------------------- /furnace/seg_opr/seg_oprs.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/6/17 上午12:43 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : seg_oprs.py 7 | from collections import OrderedDict 8 | import numpy as np 9 | 10 | import torch 11 | import torch.nn as nn 12 | 13 | 14 | class ConvBnRelu(nn.Module): 15 | def __init__(self, in_planes, out_planes, ksize, stride, pad, dilation=1, 16 | groups=1, has_bn=True, norm_layer=nn.BatchNorm2d, bn_eps=1e-5, 17 | has_relu=True, inplace=True, has_bias=False): 18 | super(ConvBnRelu, self).__init__() 19 | self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=ksize, 20 | stride=stride, padding=pad, 21 | dilation=dilation, groups=groups, bias=has_bias) 22 | self.has_bn = has_bn 23 | if self.has_bn: 24 | self.bn = norm_layer(out_planes, eps=bn_eps) 25 | self.has_relu = has_relu 26 | if self.has_relu: 27 | self.relu = nn.ReLU(inplace=inplace) 28 | 29 | def forward(self, x): 30 | x = self.conv(x) 31 | if self.has_bn: 32 | x = self.bn(x) 33 | if self.has_relu: 34 | x = self.relu(x) 35 | 36 | return x 37 | 38 | 39 | class SeparableConvBnRelu(nn.Module): 40 | def __init__(self, in_channels, out_channels, 41 | kernel_size=1, stride=1, padding=0, dilation=1, 42 | has_relu=True, norm_layer=nn.BatchNorm2d): 43 | super(SeparableConvBnRelu, self).__init__() 44 | 45 | self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, stride, 46 | padding, dilation, groups=in_channels, 47 | bias=False) 48 | self.bn = norm_layer(in_channels) 49 | self.point_wise_cbr = ConvBnRelu(in_channels, out_channels, 1, 1, 0, 50 | has_bn=True, norm_layer=norm_layer, 51 | has_relu=has_relu, has_bias=False) 52 | 53 | def forward(self, x): 54 | x = self.conv1(x) 55 | x = self.bn(x) 56 | x = self.point_wise_cbr(x) 57 | return x 58 | 59 | 60 | class GlobalAvgPool2d(nn.Module): 61 | def __init__(self): 62 | """Global average pooling over the input's spatial dimensions""" 63 | super(GlobalAvgPool2d, self).__init__() 64 | 65 | def forward(self, inputs): 66 | in_size = inputs.size() 67 | inputs = inputs.view((in_size[0], in_size[1], -1)).mean(dim=2) 68 | inputs = inputs.view(in_size[0], in_size[1], 1, 1) 69 | 70 | return inputs 71 | 72 | 73 | class SELayer(nn.Module): 74 | def __init__(self, in_planes, out_planes, reduction=16): 75 | super(SELayer, self).__init__() 76 | self.avg_pool = nn.AdaptiveAvgPool2d(1) 77 | self.fc = nn.Sequential( 78 | nn.Linear(in_planes, out_planes // reduction), 79 | nn.ReLU(inplace=True), 80 | nn.Linear(out_planes // reduction, out_planes), 81 | nn.Sigmoid() 82 | ) 83 | self.out_planes = out_planes 84 | 85 | def forward(self, x): 86 | b, c, _, _ = x.size() 87 | y = self.avg_pool(x).view(b, c) 88 | y = self.fc(y).view(b, self.out_planes, 1, 1) 89 | return y 90 | 91 | 92 | # For DFN 93 | class ChannelAttention(nn.Module): 94 | def __init__(self, in_planes, out_planes, reduction): 95 | super(ChannelAttention, self).__init__() 96 | self.channel_attention = SELayer(in_planes, out_planes, reduction) 97 | 98 | def forward(self, x1, x2): 99 | fm = torch.cat([x1, x2], 1) 100 | channel_attetion = self.channel_attention(fm) 101 | fm = x1 * channel_attetion + x2 102 | 103 | return fm 104 | 105 | 106 | class BNRefine(nn.Module): 107 | def __init__(self, in_planes, out_planes, ksize, has_bias=False, 108 | has_relu=False, norm_layer=nn.BatchNorm2d, bn_eps=1e-5): 109 | super(BNRefine, self).__init__() 110 | self.conv_bn_relu = ConvBnRelu(in_planes, out_planes, ksize, 1, 111 | ksize // 2, has_bias=has_bias, 112 | norm_layer=norm_layer, bn_eps=bn_eps) 113 | self.conv_refine = nn.Conv2d(out_planes, out_planes, kernel_size=ksize, 114 | stride=1, padding=ksize // 2, dilation=1, 115 | bias=has_bias) 116 | self.has_relu = has_relu 117 | if self.has_relu: 118 | self.relu = nn.ReLU(inplace=False) 119 | 120 | def forward(self, x): 121 | t = self.conv_bn_relu(x) 122 | t = self.conv_refine(t) 123 | if self.has_relu: 124 | return self.relu(t + x) 125 | return t + x 126 | 127 | 128 | class RefineResidual(nn.Module): 129 | def __init__(self, in_planes, out_planes, ksize, has_bias=False, 130 | has_relu=False, norm_layer=nn.BatchNorm2d, bn_eps=1e-5): 131 | super(RefineResidual, self).__init__() 132 | self.conv_1x1 = nn.Conv2d(in_planes, out_planes, kernel_size=1, 133 | stride=1, padding=0, dilation=1, 134 | bias=has_bias) 135 | self.cbr = ConvBnRelu(out_planes, out_planes, ksize, 1, 136 | ksize // 2, has_bias=has_bias, 137 | norm_layer=norm_layer, bn_eps=bn_eps) 138 | self.conv_refine = nn.Conv2d(out_planes, out_planes, kernel_size=ksize, 139 | stride=1, padding=ksize // 2, dilation=1, 140 | bias=has_bias) 141 | self.has_relu = has_relu 142 | if self.has_relu: 143 | self.relu = nn.ReLU(inplace=False) 144 | 145 | def forward(self, x): 146 | x = self.conv_1x1(x) 147 | t = self.cbr(x) 148 | t = self.conv_refine(t) 149 | if self.has_relu: 150 | return self.relu(t + x) 151 | return t + x 152 | 153 | 154 | # For BiSeNet 155 | class AttentionRefinement(nn.Module): 156 | def __init__(self, in_planes, out_planes, 157 | norm_layer=nn.BatchNorm2d): 158 | super(AttentionRefinement, self).__init__() 159 | self.conv_3x3 = ConvBnRelu(in_planes, out_planes, 3, 1, 1, 160 | has_bn=True, norm_layer=norm_layer, 161 | has_relu=True, has_bias=False) 162 | self.channel_attention = nn.Sequential( 163 | nn.AdaptiveAvgPool2d(1), 164 | ConvBnRelu(out_planes, out_planes, 1, 1, 0, 165 | has_bn=True, norm_layer=norm_layer, 166 | has_relu=False, has_bias=False), 167 | nn.Sigmoid() 168 | ) 169 | 170 | def forward(self, x): 171 | fm = self.conv_3x3(x) 172 | fm_se = self.channel_attention(fm) 173 | fm = fm * fm_se 174 | 175 | return fm 176 | 177 | 178 | class FeatureFusion(nn.Module): 179 | def __init__(self, in_planes, out_planes, 180 | reduction=1, norm_layer=nn.BatchNorm2d): 181 | super(FeatureFusion, self).__init__() 182 | self.conv_1x1 = ConvBnRelu(in_planes, out_planes, 1, 1, 0, 183 | has_bn=True, norm_layer=norm_layer, 184 | has_relu=True, has_bias=False) 185 | self.channel_attention = nn.Sequential( 186 | nn.AdaptiveAvgPool2d(1), 187 | ConvBnRelu(out_planes, out_planes // reduction, 1, 1, 0, 188 | has_bn=False, norm_layer=norm_layer, 189 | has_relu=True, has_bias=False), 190 | ConvBnRelu(out_planes // reduction, out_planes, 1, 1, 0, 191 | has_bn=False, norm_layer=norm_layer, 192 | has_relu=False, has_bias=False), 193 | nn.Sigmoid() 194 | ) 195 | 196 | def forward(self, x1, x2): 197 | fm = torch.cat([x1, x2], dim=1) 198 | fm = self.conv_1x1(fm) 199 | fm_se = self.channel_attention(fm) 200 | output = fm + fm * fm_se 201 | return output 202 | -------------------------------------------------------------------------------- /furnace/seg_opr/sgd.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/9/12 下午3:03 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : sgd.py 7 | 8 | import torch 9 | from torch.optim.sgd import SGD 10 | 11 | 12 | class StandardSGD(SGD): 13 | def step(self, closure=None): 14 | """Performs a single optimization step. 15 | Arguments: 16 | closure (callable, optional): A closure that reevaluates the model 17 | and returns the loss. 18 | """ 19 | loss = None 20 | if closure is not None: 21 | loss = closure() 22 | 23 | for group in self.param_groups: 24 | weight_decay = group['weight_decay'] 25 | momentum = group['momentum'] 26 | dampening = group['dampening'] 27 | nesterov = group['nesterov'] 28 | 29 | for p in group['params']: 30 | if p.grad is None: 31 | continue 32 | d_p = p.grad.data 33 | if weight_decay != 0: 34 | d_p.add_(weight_decay, p.data) 35 | d_p.mul_(group['lr']) 36 | if momentum != 0: 37 | param_state = self.state[p] 38 | if 'momentum_buffer' not in param_state: 39 | buf = param_state['momentum_buffer'] = torch.zeros_like( 40 | p.data) 41 | buf.mul_(momentum).add_(d_p) 42 | else: 43 | buf = param_state['momentum_buffer'] 44 | buf.mul_(momentum).add_(1 - dampening, d_p) 45 | if nesterov: 46 | d_p = d_p.add(momentum, buf) 47 | else: 48 | d_p = buf 49 | 50 | p.data.add_(-1, d_p) 51 | 52 | return loss -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/10/3 下午2:10 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : __init__.py 7 | 8 | from .syncbn import * 9 | from .parallel import * -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/comm.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/comm.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/comm.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/comm.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/comm.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/comm.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/functions.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/functions.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/functions.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/functions.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/functions.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/functions.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/parallel.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/parallel.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/parallel.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/parallel.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/parallel.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/parallel.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/syncbn.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/syncbn.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/syncbn.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/syncbn.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/__pycache__/syncbn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/__pycache__/syncbn.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/comm.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # File : comm.py 3 | # Author : Jiayuan Mao 4 | # Email : maojiayuan@gmail.com 5 | # Date : 27/01/2018 6 | # 7 | # This file is part of Synchronized-BatchNorm-PyTorch. 8 | # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch 9 | # Distributed under MIT License. 10 | 11 | import queue 12 | import collections 13 | import threading 14 | 15 | 16 | __all__ = ['FutureResult', 'SlavePipe', 'SyncMaster'] 17 | 18 | 19 | class FutureResult(object): 20 | """A thread-safe future implementation. Used only as one-to-one pipe.""" 21 | 22 | def __init__(self): 23 | self._result = None 24 | self._lock = threading.Lock() 25 | self._cond = threading.Condition(self._lock) 26 | 27 | def put(self, result): 28 | with self._lock: 29 | assert self._result is None, 'Previous result has\'t been fetched.' 30 | self._result = result 31 | self._cond.notify() 32 | 33 | def get(self): 34 | with self._lock: 35 | if self._result is None: 36 | self._cond.wait() 37 | 38 | res = self._result 39 | self._result = None 40 | return res 41 | 42 | 43 | _MasterRegistry = collections.namedtuple('MasterRegistry', ['result']) 44 | _SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result']) 45 | 46 | 47 | class SlavePipe(_SlavePipeBase): 48 | """Pipe for master-slave communication.""" 49 | 50 | def run_slave(self, msg): 51 | self.queue.put((self.identifier, msg)) 52 | ret = self.result.get() 53 | self.queue.put(True) 54 | return ret 55 | 56 | 57 | class SyncMaster(object): 58 | """An abstract `SyncMaster` object. 59 | 60 | - During the replication, as the data parallel will trigger an callback of each module, all slave devices should 61 | call `register(id)` and obtain an `SlavePipe` to communicate with the master. 62 | - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected, 63 | and passed to a registered callback. 64 | - After receiving the messages, the master device should gather the information and determine to message passed 65 | back to each slave devices. 66 | """ 67 | 68 | def __init__(self, master_callback): 69 | """ 70 | 71 | Args: 72 | master_callback: a callback to be invoked after having collected messages from slave devices. 73 | """ 74 | self._master_callback = master_callback 75 | self._queue = queue.Queue() 76 | self._registry = collections.OrderedDict() 77 | self._activated = False 78 | 79 | def register_slave(self, identifier): 80 | """ 81 | Register an slave device. 82 | 83 | Args: 84 | identifier: an identifier, usually is the device id. 85 | 86 | Returns: a `SlavePipe` object which can be used to communicate with the master device. 87 | 88 | """ 89 | if self._activated: 90 | assert self._queue.empty(), 'Queue is not clean before next initialization.' 91 | self._activated = False 92 | self._registry.clear() 93 | future = FutureResult() 94 | self._registry[identifier] = _MasterRegistry(future) 95 | return SlavePipe(identifier, self._queue, future) 96 | 97 | def run_master(self, master_msg): 98 | """ 99 | Main entry for the master device in each forward pass. 100 | The messages were first collected from each devices (including the master device), and then 101 | an callback will be invoked to compute the message to be sent back to each devices 102 | (including the master device). 103 | 104 | Args: 105 | master_msg: the message that the master want to send to itself. This will be placed as the first 106 | message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example. 107 | 108 | Returns: the message to be sent back to the master device. 109 | 110 | """ 111 | self._activated = True 112 | 113 | intermediates = [(0, master_msg)] 114 | for i in range(self.nr_slaves): 115 | intermediates.append(self._queue.get()) 116 | 117 | results = self._master_callback(intermediates) 118 | assert results[0][0] == 0, 'The first result should belongs to the master.' 119 | 120 | for i, res in results: 121 | if i == 0: 122 | continue 123 | self._registry[i].result.put(res) 124 | 125 | for i in range(self.nr_slaves): 126 | assert self._queue.get() is True 127 | 128 | return results[0][1] 129 | 130 | @property 131 | def nr_slaves(self): 132 | return len(self._registry) -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/functions.py: -------------------------------------------------------------------------------- 1 | ##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 | ## Created by: Hang Zhang 3 | ## Email: zhanghang0704@gmail.com 4 | ## Copyright (c) 2018 5 | ## 6 | ## This source code is licensed under the MIT-style license found in the 7 | ## LICENSE file in the root directory of this source tree 8 | ##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 9 | 10 | """Synchronized Cross-GPU Batch Normalization functions""" 11 | import torch 12 | from torch.autograd import Variable, Function 13 | from .src import * 14 | 15 | __all__ = ['sum_square', 'batchnormtrain'] 16 | 17 | def sum_square(input): 18 | r"""Calculate sum of elements and sum of squares for Batch Normalization""" 19 | return _sum_square.apply(input) 20 | 21 | 22 | class _sum_square(Function): 23 | @staticmethod 24 | def forward(ctx, input): 25 | ctx.save_for_backward(input) 26 | if input.is_cuda: 27 | xsum, xsqusum = gpu.sumsquare_forward(input) 28 | else: 29 | xsum, xsqusum = cpu.sumsquare_forward(input) 30 | return xsum, xsqusum 31 | 32 | @staticmethod 33 | def backward(ctx, gradSum, gradSquare): 34 | input, = ctx.saved_variables 35 | if input.is_cuda: 36 | gradInput = gpu.sumsquare_backward(input, gradSum, gradSquare) 37 | else: 38 | raise NotImplemented 39 | return gradInput 40 | 41 | 42 | class _batchnormtrain(Function): 43 | @staticmethod 44 | def forward(ctx, input, mean, std, gamma, beta): 45 | ctx.save_for_backward(input, mean, std, gamma, beta) 46 | if input.is_cuda: 47 | output = gpu.batchnorm_forward(input, mean, std, gamma, beta) 48 | else: 49 | output = cpu.batchnorm_forward(input, mean, std, gamma, beta) 50 | return output 51 | 52 | @staticmethod 53 | def backward(ctx, gradOutput): 54 | input, mean, std, gamma, beta = ctx.saved_variables 55 | if gradOutput.is_cuda: 56 | gradInput, gradMean, gradStd, gradGamma, gradBeta = \ 57 | gpu.batchnorm_backward(gradOutput, input, mean, 58 | std, gamma, beta, True) 59 | else: 60 | raise NotImplemented 61 | return gradInput, gradMean, gradStd, gradGamma, gradBeta 62 | 63 | 64 | def batchnormtrain(input, mean, std, gamma, beta): 65 | r"""Applies Batch Normalization over a 3d input that is seen as a 66 | mini-batch. 67 | 68 | .. _encoding.batchnormtrain: 69 | 70 | .. math:: 71 | 72 | y = \frac{x - \mu[x]}{ \sqrt{var[x] + \epsilon}} * \gamma + \beta 73 | 74 | Shape: 75 | - Input: :math:`(N, C)` or :math:`(N, C, L)` 76 | - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input) 77 | 78 | """ 79 | return _batchnormtrain.apply(input, mean, std, gamma, beta) 80 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/parallel.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # File : replicate.py 3 | # Author : Jiayuan Mao 4 | # Email : maojiayuan@gmail.com 5 | # Date : 27/01/2018 6 | # 7 | # This file is part of Synchronized-BatchNorm-PyTorch. 8 | # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch 9 | # Distributed under MIT License. 10 | 11 | import functools 12 | 13 | import torch 14 | from torch.nn.parallel.data_parallel import DataParallel 15 | from torch.autograd import Variable, Function 16 | import torch.cuda.comm as comm 17 | from torch.nn.parallel._functions import Broadcast 18 | 19 | # from .parallel_apply import parallel_apply 20 | 21 | torch_ver = torch.__version__[:3] 22 | 23 | __all__ = ['allreduce', 'Reduce', 'DataParallelModel', 'patch_replication_callback'] 24 | 25 | 26 | def allreduce(*inputs): 27 | """Cross GPU all reduce autograd operation for calculate mean and 28 | variance in SyncBN. 29 | """ 30 | return AllReduce.apply(*inputs) 31 | 32 | class AllReduce(Function): 33 | @staticmethod 34 | def forward(ctx, num_inputs, *inputs): 35 | ctx.num_inputs = num_inputs 36 | ctx.target_gpus = [inputs[i].get_device() for i in range(0, len(inputs), num_inputs)] 37 | inputs = [inputs[i:i + num_inputs] 38 | for i in range(0, len(inputs), num_inputs)] 39 | # sort before reduce sum 40 | inputs = sorted(inputs, key=lambda i: i[0].get_device()) 41 | results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0]) 42 | outputs = comm.broadcast_coalesced(results, ctx.target_gpus) 43 | return tuple([t for tensors in outputs for t in tensors]) 44 | 45 | @staticmethod 46 | def backward(ctx, *inputs): 47 | inputs = [i.data for i in inputs] 48 | inputs = [inputs[i:i + ctx.num_inputs] 49 | for i in range(0, len(inputs), ctx.num_inputs)] 50 | results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0]) 51 | outputs = comm.broadcast_coalesced(results, ctx.target_gpus) 52 | return (None,) + tuple([Variable(t) for tensors in outputs for t in tensors]) 53 | 54 | 55 | class Reduce(Function): 56 | @staticmethod 57 | def forward(ctx, *inputs): 58 | ctx.target_gpus = [inputs[i].get_device() for i in range(len(inputs))] 59 | inputs = sorted(inputs, key=lambda i: i.get_device()) 60 | return comm.reduce_add(inputs) 61 | 62 | @staticmethod 63 | def backward(ctx, gradOutput): 64 | return Broadcast.apply(ctx.target_gpus, gradOutput) 65 | 66 | 67 | class CallbackContext(object): 68 | pass 69 | 70 | 71 | def execute_replication_callbacks(modules): 72 | """ 73 | Execute an replication callback `__data_parallel_replicate__` on each module created by original replication. 74 | 75 | The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` 76 | 77 | Note that, as all modules are isomorphism, we assign each sub-module with a context 78 | (shared among multiple copies of this module on different devices). 79 | Through this context, different copies can share some information. 80 | 81 | We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback 82 | of any slave copies. 83 | """ 84 | master_copy = modules[0] 85 | nr_modules = len(list(master_copy.modules())) 86 | ctxs = [CallbackContext() for _ in range(nr_modules)] 87 | 88 | for i, module in enumerate(modules): 89 | for j, m in enumerate(module.modules()): 90 | if hasattr(m, '__data_parallel_replicate__'): 91 | m.__data_parallel_replicate__(ctxs[j], i) 92 | 93 | 94 | 95 | class DataParallelModel(DataParallel): 96 | """Implements data parallelism at the module level. 97 | 98 | This container parallelizes the application of the given module by 99 | splitting the input across the specified devices by chunking in the 100 | batch dimension. 101 | In the forward pass, the module is replicated on each device, 102 | and each replica handles a portion of the input. During the backwards pass, gradients from each replica are summed into the original module. 103 | Note that the outputs are not gathered, please use compatible 104 | :class:`encoding.parallel.DataParallelCriterion`. 105 | 106 | The batch size should be larger than the number of GPUs used. It should 107 | also be an integer multiple of the number of GPUs so that each chunk is 108 | the same size (so that each GPU processes the same number of samples). 109 | 110 | Args: 111 | module: module to be parallelized 112 | device_ids: CUDA devices (default: all devices) 113 | 114 | Reference: 115 | Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi, 116 | Amit Agrawal. “Context Encoding for Semantic Segmentation. 117 | *The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018* 118 | 119 | Example:: 120 | 121 | >>> net = DataParallelModel(model, device_ids=[0, 1, 2]) 122 | >>> y = net(x) 123 | """ 124 | def gather(self, outputs, output_device): 125 | return outputs 126 | 127 | def replicate(self, module, device_ids): 128 | modules = super(DataParallelModel, self).replicate(module, device_ids) 129 | execute_replication_callbacks(modules) 130 | return modules 131 | 132 | # def parallel_apply(self, replicas, inputs, kwargs): 133 | # return parallel_apply(replicas, inputs, kwargs, 134 | # self.device_ids[:len(replicas)]) 135 | 136 | 137 | def patch_replication_callback(data_parallel): 138 | """ 139 | Monkey-patch an existing `DataParallel` object. Add the replication callback. 140 | Useful when you have customized `DataParallel` implementation. 141 | 142 | Examples: 143 | > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) 144 | > sync_bn = DataParallel(sync_bn, device_ids=[0, 1]) 145 | > patch_replication_callback(sync_bn) 146 | # this is equivalent to 147 | > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) 148 | > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) 149 | """ 150 | 151 | assert isinstance(data_parallel, DataParallel) 152 | 153 | old_replicate = data_parallel.replicate 154 | 155 | @functools.wraps(old_replicate) 156 | def new_replicate(module, device_ids): 157 | modules = old_replicate(module, device_ids) 158 | execute_replication_callbacks(modules) 159 | return modules 160 | 161 | data_parallel.replicate = new_replicate 162 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/parallel_apply.py: -------------------------------------------------------------------------------- 1 | # import threading 2 | import queue 3 | import torch 4 | import torch.multiprocessing as mp 5 | # from pathos.multiprocessing import ProcessPool as Pool 6 | from torch.cuda._utils import _get_device_index 7 | 8 | 9 | def get_a_var(obj): 10 | if isinstance(obj, torch.Tensor): 11 | return obj 12 | 13 | if isinstance(obj, list) or isinstance(obj, tuple): 14 | for result in map(get_a_var, obj): 15 | if isinstance(result, torch.Tensor): 16 | return result 17 | if isinstance(obj, dict): 18 | for result in map(get_a_var, obj.items()): 19 | if isinstance(result, torch.Tensor): 20 | return result 21 | return None 22 | 23 | 24 | def parallel_apply(modules, inputs, kwargs_tup=None, devices=None): 25 | r"""Applies each `module` in :attr:`modules` in parallel on arguments 26 | contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword) 27 | on each of :attr:`devices`. 28 | Args: 29 | modules (Module): modules to be parallelized 30 | inputs (tensor): inputs to the modules 31 | devices (list of int or torch.device): CUDA devices 32 | :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and 33 | :attr:`devices` (if given) should all have same length. Moreover, each 34 | element of :attr:`inputs` can either be a single object as the only argument 35 | to a module, or a collection of positional arguments. 36 | """ 37 | assert len(modules) == len(inputs) 38 | if kwargs_tup is not None: 39 | assert len(modules) == len(kwargs_tup) 40 | else: 41 | kwargs_tup = ({},) * len(modules) 42 | if devices is not None: 43 | assert len(modules) == len(devices) 44 | else: 45 | devices = [None] * len(modules) 46 | devices = list(map(lambda x: _get_device_index(x, True), devices)) 47 | context = mp.get_context('spawn') 48 | # lock = threading.Lock() 49 | # results = {} 50 | # results = [] 51 | # pool = context.Pool(len(devices)) 52 | results_queue = queue.Queue(len(devices)) 53 | grad_enabled = torch.is_grad_enabled() 54 | 55 | def _worker(module, input, kwargs, device=None): 56 | torch.set_grad_enabled(grad_enabled) 57 | if device is None: 58 | device = get_a_var(input).get_device() 59 | try: 60 | with torch.cuda.device(device): 61 | # this also avoids accidental slicing of `input` if it is a Tensor 62 | if not isinstance(input, (list, tuple)): 63 | input = (input,) 64 | output = module(*input, **kwargs) 65 | results_queue.put(output) 66 | # with lock: 67 | # results[i] = output 68 | except Exception as e: 69 | results_queue.put(e) 70 | # with lock: 71 | # results[i] = e 72 | 73 | if len(modules) > 1: 74 | # pool.map(_worker, [modules, inputs, kwargs_tup, devices]) 75 | processes = [context.Process(target=_worker, 76 | args=(i, module, input, kwargs, device)) 77 | for i, (module, input, kwargs, device) in 78 | enumerate(zip(modules, inputs, kwargs_tup, devices))] 79 | 80 | for process in processes: 81 | process.start() 82 | for process in processes: 83 | process.join() 84 | else: 85 | _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0]) 86 | 87 | outputs = [] 88 | for i in range(len(inputs)): 89 | output = results_queue.get() 90 | if isinstance(output, Exception): 91 | raise output 92 | outputs.append(output) 93 | return outputs 94 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from torch.utils.cpp_extension import load 4 | 5 | cwd = os.path.dirname(os.path.realpath(__file__)) 6 | cpu_path = os.path.join(cwd, 'cpu') 7 | gpu_path = os.path.join(cwd, 'gpu') 8 | 9 | cpu = load('syncbn_cpu', [ 10 | os.path.join(cpu_path, 'operator.cpp'), 11 | os.path.join(cpu_path, 'syncbn_cpu.cpp'), 12 | ], build_directory=cpu_path, verbose=False) 13 | 14 | if torch.cuda.is_available(): 15 | gpu = load('syncbn_gpu', [ 16 | os.path.join(gpu_path, 'operator.cpp'), 17 | os.path.join(gpu_path, 'syncbn_kernel.cu'), 18 | ], build_directory=gpu_path, verbose=False) 19 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/.ninja_deps: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/cpu/.ninja_deps -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/.ninja_log: -------------------------------------------------------------------------------- 1 | # ninja log v5 2 | 1 3006 1563332513 syncbn_cpu.o 486ee2c6335a262c 3 | 1 6096 1563332516 operator.o df1c06f439a829e3 4 | 6096 6262 1563332517 syncbn_cpu.so 7b7138baea8e4fe0 5 | 0 3376 1578576757073544196 syncbn_cpu.o 238aaa649062d1c 6 | 0 4373 1578576758073181846 operator.o eedcce4cadeab94a 7 | 4373 4493 1578576758193138364 syncbn_cpu.so 7b7138baea8e4fe0 8 | 30 4937 1617328371000000000 syncbn_cpu.o 9e462707afc905b4 9 | 30 9254 1617328375000000000 operator.o 38caa69f5662504d 10 | 9254 9523 1617328375000000000 syncbn_cpu.so f2ec696724cd4854 11 | 1 2576 1617337607 syncbn_cpu.o 12505825316ac744 12 | 1 5784 1617337610 operator.o 44a849fcf5a36c1f 13 | 5784 5901 1617337610 syncbn_cpu.so 7b7138baea8e4fe0 14 | 0 3924 1617338218000000000 syncbn_cpu.o 9e462707afc905b4 15 | 0 8972 1617338223000000000 operator.o 38caa69f5662504d 16 | 8972 9227 1617338223000000000 syncbn_cpu.so f2ec696724cd4854 17 | 1 2824 1617344917 syncbn_cpu.o 12505825316ac744 18 | 1 5115 1617344920 operator.o 44a849fcf5a36c1f 19 | 5115 5235 1617344920 syncbn_cpu.so 7b7138baea8e4fe0 20 | 1 3927 1617345308000000000 syncbn_cpu.o 9e462707afc905b4 21 | 1 8535 1617345313000000000 operator.o 38caa69f5662504d 22 | 8535 8807 1617345313000000000 syncbn_cpu.so f2ec696724cd4854 23 | 1 3809 1617346702000000000 syncbn_cpu.o cead0f3b19ad107 24 | 1 8661 1617346706000000000 operator.o 3c1189354cdcedbd 25 | 8661 8896 1617346707000000000 syncbn_cpu.so 96100cc8653a7565 26 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/cpu/__init__.py -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/build.ninja: -------------------------------------------------------------------------------- 1 | ninja_required_version = 1.3 2 | cxx = c++ 3 | 4 | cflags = -DTORCH_EXTENSION_NAME=syncbn_cpu -DTORCH_API_INCLUDE_EXTENSION_H -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include/TH -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include/THC -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/include/python3.7m -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14 5 | post_cflags = 6 | ldflags = -shared -L/mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/lib -lc10 -ltorch_cpu -ltorch -ltorch_python 7 | 8 | rule compile 9 | command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags 10 | depfile = $out.d 11 | deps = gcc 12 | 13 | rule link 14 | command = $cxx $in $ldflags -o $out 15 | 16 | build operator.o: compile /mnt/lustre/liushinan/cyj/start/furnace/seg_opr/sync_bn/src/cpu/operator.cpp 17 | build syncbn_cpu.o: compile /mnt/lustre/liushinan/cyj/start/furnace/seg_opr/sync_bn/src/cpu/syncbn_cpu.cpp 18 | 19 | build syncbn_cpu.so: link operator.o syncbn_cpu.o 20 | 21 | default syncbn_cpu.so 22 | 23 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/dist/syncbn_cpu-0.0.0-py3.6-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/cpu/dist/syncbn_cpu-0.0.0-py3.6-linux-x86_64.egg -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/operator.cpp: -------------------------------------------------------------------------------- 1 | #include "operator.h" 2 | 3 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 4 | m.def("batchnorm_forward", &BatchNorm_Forward_CPU, "BatchNorm forward (CPU)"); 5 | m.def("batchnorm_backward", &BatchNorm_Backward_CPU, "BatchNorm backward (CPU)"); 6 | m.def("sumsquare_forward", &Sum_Square_Forward_CPU, "SumSqu forward (CPU)"); 7 | m.def("sumsquare_backward", &Sum_Square_Backward_CPU, "SumSqu backward (CPU)"); 8 | } 9 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/operator.h: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | at::Tensor BatchNorm_Forward_CPU( 5 | const at::Tensor input_, 6 | const at::Tensor mean_, 7 | const at::Tensor std_, 8 | const at::Tensor gamma_, 9 | const at::Tensor beta_); 10 | 11 | std::vector BatchNorm_Backward_CPU( 12 | const at::Tensor gradoutput_, 13 | const at::Tensor input_, 14 | const at::Tensor mean_, 15 | const at::Tensor std_, 16 | const at::Tensor gamma_, 17 | const at::Tensor beta_, 18 | bool train); 19 | 20 | std::vector Sum_Square_Forward_CPU( 21 | const at::Tensor input_); 22 | 23 | at::Tensor Sum_Square_Backward_CPU( 24 | const at::Tensor input_, 25 | const at::Tensor gradSum_, 26 | const at::Tensor gradSquare_); -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/operator.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/cpu/operator.o -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CppExtension 3 | 4 | setup( 5 | name='syncbn_cpu', 6 | ext_modules=[ 7 | CppExtension('syncbn_cpu', [ 8 | 'operator.cpp', 9 | 'syncbn_cpu.cpp', 10 | ]), 11 | ], 12 | cmdclass={ 13 | 'build_ext': BuildExtension 14 | }) 15 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/syncbn_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | at::Tensor broadcast_to(at::Tensor v, at::Tensor x) { 6 | if (x.ndimension() == 2) { 7 | return v; 8 | } else { 9 | std::vector broadcast_size = {1, -1}; 10 | for (int64_t i = 2; i < x.ndimension(); ++i) 11 | broadcast_size.push_back(1); 12 | 13 | return v.view(broadcast_size); 14 | } 15 | } 16 | 17 | at::Tensor BatchNorm_Forward_CPU( 18 | const at::Tensor input, 19 | const at::Tensor mean, 20 | const at::Tensor std, 21 | const at::Tensor gamma, 22 | const at::Tensor beta) { 23 | auto output = (input - broadcast_to(mean, input)) / broadcast_to(std, input); 24 | output = output * broadcast_to(gamma, input) + broadcast_to(beta, input); 25 | return output; 26 | } 27 | 28 | // Not implementing CPU backward for now 29 | std::vector BatchNorm_Backward_CPU( 30 | const at::Tensor gradoutput, 31 | const at::Tensor input, 32 | const at::Tensor mean, 33 | const at::Tensor std, 34 | const at::Tensor gamma, 35 | const at::Tensor beta, 36 | bool train) { 37 | /* outputs*/ 38 | at::Tensor gradinput = at::zeros_like(input); 39 | at::Tensor gradgamma = at::zeros_like(gamma); 40 | at::Tensor gradbeta = at::zeros_like(beta); 41 | at::Tensor gradMean = at::zeros_like(mean); 42 | at::Tensor gradStd = at::zeros_like(std); 43 | return {gradinput, gradMean, gradStd, gradgamma, gradbeta}; 44 | } 45 | 46 | std::vector Sum_Square_Forward_CPU( 47 | const at::Tensor input) { 48 | /* outputs */ 49 | at::Tensor sum = torch::zeros({input.size(1)}, input.options()); 50 | at::Tensor square = torch::zeros({input.size(1)}, input.options()); 51 | return {sum, square}; 52 | } 53 | 54 | at::Tensor Sum_Square_Backward_CPU( 55 | const at::Tensor input, 56 | const at::Tensor gradSum, 57 | const at::Tensor gradSquare) { 58 | /* outputs */ 59 | at::Tensor gradInput = at::zeros_like(input); 60 | return gradInput; 61 | } 62 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/syncbn_cpu.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: syncbn-cpu 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/syncbn_cpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/cpu/syncbn_cpu.o -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/cpu/syncbn_cpu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/cpu/syncbn_cpu.so -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/.ninja_deps: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/gpu/.ninja_deps -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/.ninja_log: -------------------------------------------------------------------------------- 1 | # ninja log v5 2 | 1 6029 1563332523 operator.o 93fbaee254d44db4 3 | 1 11088 1563332528 syncbn_kernel.cuda.o ec50d81437939f2c 4 | 11088 11258 1563332528 syncbn_gpu.so a2b728e60c853ec3 5 | 0 2904 1578576761208045520 operator.o cecba1516d789115 6 | 0 7039 1578576765338548294 syncbn_kernel.cuda.o a19378e9f1e5d587 7 | 7039 7134 1578576765434513509 syncbn_gpu.so a2b728e60c853ec3 8 | 1 8148 1617328383000000000 operator.o 4baea6eee61c764a 9 | 1 22462 1617328398000000000 syncbn_kernel.cuda.o e2a2dae82774a21e 10 | 22463 22801 1617328398000000000 syncbn_gpu.so 8be6fcd47ba2b271 11 | 0 4933 1617337615 operator.o 4c1e266b352b1f39 12 | 0 8131 1617337618 syncbn_kernel.cuda.o abcbcfa4a428e6b3 13 | 8131 8233 1617337618 syncbn_gpu.so a2b728e60c853ec3 14 | 5 8183 1617338231000000000 operator.o 4baea6eee61c764a 15 | 5 13959 1617338237000000000 syncbn_kernel.cuda.o e2a2dae82774a21e 16 | 13959 14198 1617338237000000000 syncbn_gpu.so 8be6fcd47ba2b271 17 | 1 4907 1617344925 operator.o 4c1e266b352b1f39 18 | 1 7768 1617344928 syncbn_kernel.cuda.o abcbcfa4a428e6b3 19 | 7768 7874 1617344928 syncbn_gpu.so a2b728e60c853ec3 20 | 21 8085 1617345321000000000 operator.o 4baea6eee61c764a 21 | 21 13693 1617345327000000000 syncbn_kernel.cuda.o e2a2dae82774a21e 22 | 13698 13960 1617345327000000000 syncbn_gpu.so 8be6fcd47ba2b271 23 | 4 7772 1617346715000000000 operator.o cff9f37d81fb290d 24 | 5 12976 1617346720000000000 syncbn_kernel.cuda.o 7649931171c72623 25 | 12976 13237 1617346720000000000 syncbn_gpu.so 6c2fe5233eb33b0e 26 | 1 18316 1617351381000000000 operator.o 8b96510ead39db99 27 | 1 21999 1617351385000000000 syncbn_kernel.cuda.o 581d7ea8c9f6e90d 28 | 21999 22404 1617351385000000000 syncbn_gpu.so 4bbd9457080b3d5f 29 | 1 8208 1617351740000000000 operator.o cff9f37d81fb290d 30 | 1 12846 1617351744000000000 syncbn_kernel.cuda.o 7649931171c72623 31 | 12846 13134 1617351744000000000 syncbn_gpu.so 6c2fe5233eb33b0e 32 | 1 12669 1617351769000000000 operator.o 8b96510ead39db99 33 | 2 17252 1617351773000000000 syncbn_kernel.cuda.o 581d7ea8c9f6e90d 34 | 17252 17525 1617351774000000000 syncbn_gpu.so 4bbd9457080b3d5f 35 | 0 10940 1617352935000000000 operator.o cff9f37d81fb290d 36 | 0 17323 1617352941000000000 syncbn_kernel.cuda.o 7649931171c72623 37 | 17323 17932 1617352942000000000 syncbn_gpu.so 6c2fe5233eb33b0e 38 | 0 16416 1617353014000000000 operator.o 8b96510ead39db99 39 | 0 22125 1617353019000000000 syncbn_kernel.cuda.o 581d7ea8c9f6e90d 40 | 22125 22424 1617353020000000000 syncbn_gpu.so 4bbd9457080b3d5f 41 | 1 9054 1617353253000000000 operator.o cff9f37d81fb290d 42 | 1 13386 1617353257000000000 syncbn_kernel.cuda.o 7649931171c72623 43 | 13392 13676 1617353258000000000 syncbn_gpu.so 6c2fe5233eb33b0e 44 | 1 13591 1617353313000000000 operator.o 8b96510ead39db99 45 | 1 17494 1617353317000000000 syncbn_kernel.cuda.o 581d7ea8c9f6e90d 46 | 17494 17767 1617353317000000000 syncbn_gpu.so 4bbd9457080b3d5f 47 | 1 8158 1617353350000000000 operator.o cff9f37d81fb290d 48 | 1 12504 1617353355000000000 syncbn_kernel.cuda.o 7649931171c72623 49 | 12504 12806 1617353355000000000 syncbn_gpu.so 6c2fe5233eb33b0e 50 | 5 13534 1617353515000000000 operator.o 8b96510ead39db99 51 | 5 18083 1617353520000000000 syncbn_kernel.cuda.o 581d7ea8c9f6e90d 52 | 18083 18375 1617353520000000000 syncbn_gpu.so 4bbd9457080b3d5f 53 | 1 8932 1617353728000000000 operator.o cff9f37d81fb290d 54 | 1 12531 1617353731000000000 syncbn_kernel.cuda.o 7649931171c72623 55 | 12532 12826 1617353732000000000 syncbn_gpu.so 6c2fe5233eb33b0e 56 | 1 8570 1617368450000000000 operator.o a709041194a42ccf 57 | 1 20269 1617368462000000000 syncbn_kernel.cuda.o d470b9775a8606af 58 | 20270 20582 1617368462000000000 syncbn_gpu.so 73957be52d8f19b6 59 | 1 8643 1617369090000000000 operator.o cff9f37d81fb290d 60 | 1 13121 1617369094000000000 syncbn_kernel.cuda.o 7649931171c72623 61 | 13121 13430 1617369094000000000 syncbn_gpu.so 6c2fe5233eb33b0e 62 | 21 12357 1617418761000000000 operator.o a709041194a42ccf 63 | 22 16395 1617418765000000000 syncbn_kernel.cuda.o d470b9775a8606af 64 | 16396 16691 1617418765000000000 syncbn_gpu.so 73957be52d8f19b6 65 | 1 14122 1617525702000000000 operator.o cff9f37d81fb290d 66 | 1 18475 1617525706000000000 syncbn_kernel.cuda.o 7649931171c72623 67 | 18475 18974 1617525707000000000 syncbn_gpu.so 6c2fe5233eb33b0e 68 | 10 11850 1617699608000000000 operator.o a709041194a42ccf 69 | 10 14614 1617699611000000000 syncbn_kernel.cuda.o d470b9775a8606af 70 | 14620 14985 1617699611000000000 syncbn_gpu.so 73957be52d8f19b6 71 | 82 9491 1617702181000000000 operator.o cff9f37d81fb290d 72 | 82 14360 1617702186000000000 syncbn_kernel.cuda.o 7649931171c72623 73 | 14360 14651 1617702186000000000 syncbn_gpu.so 6c2fe5233eb33b0e 74 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/gpu/__init__.py -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/build.ninja: -------------------------------------------------------------------------------- 1 | ninja_required_version = 1.3 2 | cxx = c++ 3 | nvcc = /mnt/lustre/share/cuda-10.1/bin/nvcc 4 | 5 | cflags = -DTORCH_EXTENSION_NAME=syncbn_gpu -DTORCH_API_INCLUDE_EXTENSION_H -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include/TH -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include/THC -isystem /mnt/lustre/share/cuda-10.1/include -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/include/python3.7m -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14 6 | post_cflags = 7 | cuda_cflags = -DTORCH_EXTENSION_NAME=syncbn_gpu -DTORCH_API_INCLUDE_EXTENSION_H -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include/TH -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/include/THC -isystem /mnt/lustre/share/cuda-10.1/include -isystem /mnt/lustre/liushinan/miniconda3/envs/t16/include/python3.7m -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_70,code=sm_70 --compiler-options '-fPIC' -std=c++14 8 | cuda_post_cflags = 9 | ldflags = -shared -L/mnt/lustre/liushinan/miniconda3/envs/t16/lib/python3.7/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/mnt/lustre/share/cuda-10.1/lib64 -lcudart 10 | 11 | rule compile 12 | command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags 13 | depfile = $out.d 14 | deps = gcc 15 | 16 | rule cuda_compile 17 | command = $nvcc $cuda_cflags -c $in -o $out $cuda_post_cflags 18 | 19 | rule link 20 | command = $cxx $in $ldflags -o $out 21 | 22 | build operator.o: compile /mnt/lustre/liushinan/cyj/start/furnace/seg_opr/sync_bn/src/gpu/operator.cpp 23 | build syncbn_kernel.cuda.o: cuda_compile /mnt/lustre/liushinan/cyj/start/furnace/seg_opr/sync_bn/src/gpu/syncbn_kernel.cu 24 | 25 | build syncbn_gpu.so: link operator.o syncbn_kernel.cuda.o 26 | 27 | default syncbn_gpu.so 28 | 29 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/common.h: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | static const unsigned WARP_SIZE = 32; 5 | 6 | // The maximum number of threads in a block 7 | static const unsigned MAX_BLOCK_SIZE = 512U; 8 | 9 | template 10 | struct ScalarConvert { 11 | static __host__ __device__ __forceinline__ Out to(const In v) { return (Out) v; } 12 | }; 13 | 14 | // Number of threads in a block given an input size up to MAX_BLOCK_SIZE 15 | static int getNumThreads(int nElem) { 16 | int threadSizes[5] = { 32, 64, 128, 256, MAX_BLOCK_SIZE }; 17 | for (int i = 0; i != 5; ++i) { 18 | if (nElem <= threadSizes[i]) { 19 | return threadSizes[i]; 20 | } 21 | } 22 | return MAX_BLOCK_SIZE; 23 | } 24 | 25 | // Returns the index of the most significant 1 bit in `val`. 26 | __device__ __forceinline__ int getMSB(int val) { 27 | return 31 - __clz(val); 28 | } 29 | 30 | template 31 | __device__ __forceinline__ T WARP_SHFL_XOR(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) 32 | { 33 | #if CUDA_VERSION >= 9000 34 | return __shfl_xor_sync(mask, value, laneMask, width); 35 | #else 36 | return __shfl_xor(value, laneMask, width); 37 | #endif 38 | } 39 | 40 | // Sum across all threads within a warp 41 | template 42 | static __device__ __forceinline__ T warpSum(T val) { 43 | #if __CUDA_ARCH__ >= 300 44 | for (int i = 0; i < getMSB(WARP_SIZE); ++i) { 45 | val += WARP_SHFL_XOR(val, 1 << i, WARP_SIZE); 46 | } 47 | #else 48 | __shared__ T values[MAX_BLOCK_SIZE]; 49 | values[threadIdx.x] = val; 50 | __threadfence_block(); 51 | const int base = (threadIdx.x / WARP_SIZE) * WARP_SIZE; 52 | for (int i = 1; i < WARP_SIZE; i++) { 53 | val += values[base + ((i + threadIdx.x) % WARP_SIZE)]; 54 | } 55 | #endif 56 | return val; 57 | } 58 | 59 | template 60 | struct Float2 { 61 | Acctype v1, v2; 62 | __device__ Float2() {} 63 | __device__ Float2(DType v1, DType v2) : v1(ScalarConvert::to(v1)), v2(ScalarConvert::to(v2)) {} 64 | __device__ Float2(DType v) : v1(ScalarConvert::to(v)), v2(ScalarConvert::to(v)) {} 65 | __device__ Float2(int v) : v1(ScalarConvert::to(v)), v2(ScalarConvert::to(v)) {} 66 | __device__ Float2& operator+=(const Float2& a) { 67 | v1 += a.v1; 68 | v2 += a.v2; 69 | return *this; 70 | } 71 | }; 72 | 73 | template 74 | static __device__ __forceinline__ Float2 warpSum(Float2 value) { 75 | value.v1 = warpSum(value.v1); 76 | value.v2 = warpSum(value.v2); 77 | return value; 78 | } 79 | 80 | template 81 | __device__ T reduceD( 82 | Op op, int b, int i, int k, int D) { 83 | T sum = 0; 84 | for (int x = threadIdx.x; x < D; x += blockDim.x) { 85 | sum += op(b,i,k,x); 86 | } 87 | // sum over NumThreads within a warp 88 | sum = warpSum(sum); 89 | 90 | // 'transpose', and reduce within warp again 91 | __shared__ T shared[32]; 92 | 93 | __syncthreads(); 94 | if (threadIdx.x % WARP_SIZE == 0) { 95 | if (threadIdx.x / WARP_SIZE < 32) { 96 | shared[threadIdx.x / WARP_SIZE] = sum; 97 | } 98 | } 99 | if (threadIdx.x >= blockDim.x / WARP_SIZE && threadIdx.x < WARP_SIZE) { 100 | // zero out the other entries in shared 101 | shared[threadIdx.x] = (T) 0; 102 | } 103 | __syncthreads(); 104 | if (threadIdx.x / WARP_SIZE == 0) { 105 | sum = warpSum(shared[threadIdx.x]); 106 | if (threadIdx.x == 0) { 107 | shared[0] = sum; 108 | } 109 | } 110 | __syncthreads(); 111 | 112 | // Everyone picks it up, should be broadcast into the whole gradInput 113 | return shared[0]; 114 | } 115 | 116 | template 117 | __device__ T reduceN( 118 | Op op, int b, int k, int d, int N) { 119 | T sum = 0; 120 | for (int x = threadIdx.x; x < N; x += blockDim.x) { 121 | sum += op(b,x,k,d); 122 | } 123 | // sum over NumThreads within a warp 124 | sum = warpSum(sum); 125 | 126 | // 'transpose', and reduce within warp again 127 | __shared__ T shared[32]; 128 | 129 | __syncthreads(); 130 | if (threadIdx.x % WARP_SIZE == 0) { 131 | if (threadIdx.x / WARP_SIZE < 32) { 132 | shared[threadIdx.x / WARP_SIZE] = sum; 133 | } 134 | } 135 | if (threadIdx.x >= blockDim.x / WARP_SIZE && threadIdx.x < WARP_SIZE) { 136 | // zero out the other entries in shared 137 | shared[threadIdx.x] = (T) 0; 138 | } 139 | __syncthreads(); 140 | if (threadIdx.x / WARP_SIZE == 0) { 141 | sum = warpSum(shared[threadIdx.x]); 142 | if (threadIdx.x == 0) { 143 | shared[0] = sum; 144 | } 145 | } 146 | __syncthreads(); 147 | 148 | // Everyone picks it up, should be broadcast into the whole gradInput 149 | return shared[0]; 150 | } 151 | 152 | template 153 | __device__ T reduceK( 154 | Op op, int b, int i, int d, int K) { 155 | T sum = 0; 156 | for (int x = threadIdx.x; x < K; x += blockDim.x) { 157 | sum += op(b,i,x,d); 158 | } 159 | // sum over NumThreads within a warp 160 | sum = warpSum(sum); 161 | 162 | // 'transpose', and reduce within warp again 163 | __shared__ T shared[32]; 164 | 165 | __syncthreads(); 166 | if (threadIdx.x % WARP_SIZE == 0) { 167 | if (threadIdx.x / WARP_SIZE < 32) { 168 | shared[threadIdx.x / WARP_SIZE] = sum; 169 | } 170 | } 171 | if (threadIdx.x >= blockDim.x / WARP_SIZE && threadIdx.x < WARP_SIZE) { 172 | // zero out the other entries in shared 173 | shared[threadIdx.x] = (T) 0; 174 | } 175 | __syncthreads(); 176 | if (threadIdx.x / WARP_SIZE == 0) { 177 | sum = warpSum(shared[threadIdx.x]); 178 | if (threadIdx.x == 0) { 179 | shared[0] = sum; 180 | } 181 | } 182 | __syncthreads(); 183 | 184 | // Everyone picks it up, should be broadcast into the whole gradInput 185 | return shared[0]; 186 | } 187 | 188 | template 189 | __device__ T reduceBN( 190 | Op op, 191 | int k, int d, int B, int N) { 192 | T sum = 0; 193 | for (int batch = 0; batch < B; ++batch) { 194 | for (int x = threadIdx.x; x < N; x += blockDim.x) { 195 | sum += op(batch,x,k,d); 196 | } 197 | } 198 | // sum over NumThreads within a warp 199 | sum = warpSum(sum); 200 | // 'transpose', and reduce within warp again 201 | __shared__ T shared[32]; 202 | 203 | __syncthreads(); 204 | if (threadIdx.x % WARP_SIZE == 0) { 205 | if (threadIdx.x / WARP_SIZE < 32) { 206 | shared[threadIdx.x / WARP_SIZE] = sum; 207 | } 208 | } 209 | if (threadIdx.x >= blockDim.x / WARP_SIZE && threadIdx.x < WARP_SIZE) { 210 | // zero out the other entries in shared 211 | shared[threadIdx.x] = (T) 0; 212 | } 213 | __syncthreads(); 214 | if (threadIdx.x / WARP_SIZE == 0) { 215 | sum = warpSum(shared[threadIdx.x]); 216 | if (threadIdx.x == 0) { 217 | shared[0] = sum; 218 | } 219 | } 220 | __syncthreads(); 221 | 222 | // Everyone picks it up, should be broadcast into the whole gradInput 223 | return shared[0]; 224 | } 225 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/device_tensor.h: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | template 4 | struct DeviceTensor { 5 | public: 6 | inline __device__ __host__ DeviceTensor(DType *p, const int *size) 7 | : dptr_(p) { 8 | for (int i = 0; i < Dim; ++i) { 9 | size_[i] = size ? size[i] : 0; 10 | } 11 | } 12 | 13 | inline __device__ __host__ unsigned getSize(const int i) const { 14 | assert(i < Dim); 15 | return size_[i]; 16 | } 17 | 18 | inline __device__ __host__ int numElements() const { 19 | int n = 1; 20 | for (int i = 0; i < Dim; ++i) { 21 | n *= size_[i]; 22 | } 23 | return n; 24 | } 25 | 26 | inline __device__ __host__ DeviceTensor select(const size_t x) const { 27 | assert(Dim > 1); 28 | int offset = x; 29 | for (int i = 1; i < Dim; ++i) { 30 | offset *= size_[i]; 31 | } 32 | DeviceTensor tensor(dptr_ + offset, nullptr); 33 | for (int i = 0; i < Dim - 1; ++i) { 34 | tensor.size_[i] = this->size_[i+1]; 35 | } 36 | return tensor; 37 | } 38 | 39 | inline __device__ __host__ DeviceTensor operator[](const size_t x) const { 40 | assert(Dim > 1); 41 | int offset = x; 42 | for (int i = 1; i < Dim; ++i) { 43 | offset *= size_[i]; 44 | } 45 | DeviceTensor tensor(dptr_ + offset, nullptr); 46 | for (int i = 0; i < Dim - 1; ++i) { 47 | tensor.size_[i] = this->size_[i+1]; 48 | } 49 | return tensor; 50 | } 51 | 52 | inline __device__ __host__ size_t InnerSize() const { 53 | assert(Dim >= 3); 54 | size_t sz = 1; 55 | for (size_t i = 2; i < Dim; ++i) { 56 | sz *= size_[i]; 57 | } 58 | return sz; 59 | } 60 | 61 | inline __device__ __host__ size_t ChannelCount() const { 62 | assert(Dim >= 3); 63 | return size_[1]; 64 | } 65 | 66 | inline __device__ __host__ DType* data_ptr() const { 67 | return dptr_; 68 | } 69 | 70 | DType *dptr_; 71 | int size_[Dim]; 72 | }; 73 | 74 | template 75 | struct DeviceTensor { 76 | inline __device__ __host__ DeviceTensor(DType *p, const int *size) 77 | : dptr_(p) { 78 | size_[0] = size ? size[0] : 0; 79 | } 80 | 81 | inline __device__ __host__ unsigned getSize(const int i) const { 82 | assert(i == 0); 83 | return size_[0]; 84 | } 85 | 86 | inline __device__ __host__ int numElements() const { 87 | return size_[0]; 88 | } 89 | 90 | inline __device__ __host__ DType &operator[](const size_t x) const { 91 | return *(dptr_ + x); 92 | } 93 | 94 | inline __device__ __host__ DType* data_ptr() const { 95 | return dptr_; 96 | } 97 | 98 | DType *dptr_; 99 | int size_[1]; 100 | }; 101 | 102 | template 103 | static DeviceTensor devicetensor(const at::Tensor &blob) { 104 | DType *data = blob.data(); 105 | DeviceTensor tensor(data, nullptr); 106 | for (int i = 0; i < Dim; ++i) { 107 | tensor.size_[i] = blob.size(i); 108 | } 109 | return tensor; 110 | } 111 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/dist/syncbn_gpu-0.0.0-py3.6-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/gpu/dist/syncbn_gpu-0.0.0-py3.6-linux-x86_64.egg -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/operator.cpp: -------------------------------------------------------------------------------- 1 | #include "operator.h" 2 | 3 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 4 | m.def("batchnorm_forward", &BatchNorm_Forward_CUDA, "BatchNorm forward (CUDA)"); 5 | m.def("batchnorm_backward", &BatchNorm_Backward_CUDA, "BatchNorm backward (CUDA)"); 6 | m.def("sumsquare_forward", &Sum_Square_Forward_CUDA, "SumSqu forward (CUDA)"); 7 | m.def("sumsquare_backward", &Sum_Square_Backward_CUDA, "SumSqu backward (CUDA)"); 8 | 9 | } 10 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/operator.h: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | 5 | at::Tensor BatchNorm_Forward_CUDA( 6 | const at::Tensor input_, 7 | const at::Tensor mean_, 8 | const at::Tensor std_, 9 | const at::Tensor gamma_, 10 | const at::Tensor beta_); 11 | 12 | std::vector BatchNorm_Backward_CUDA( 13 | const at::Tensor gradoutput_, 14 | const at::Tensor input_, 15 | const at::Tensor mean_, 16 | const at::Tensor std_, 17 | const at::Tensor gamma_, 18 | const at::Tensor beta_, 19 | bool train); 20 | 21 | std::vector Sum_Square_Forward_CUDA( 22 | const at::Tensor input_); 23 | 24 | at::Tensor Sum_Square_Backward_CUDA( 25 | const at::Tensor input_, 26 | const at::Tensor gradSum_, 27 | const at::Tensor gradSquare_); 28 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/operator.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/gpu/operator.o -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | setup( 5 | name='syncbn_gpu', 6 | ext_modules=[ 7 | CUDAExtension('syncbn_gpu', [ 8 | 'operator.cpp', 9 | 'syncbn_kernel.cu', 10 | ]), 11 | ], 12 | cmdclass={ 13 | 'build_ext': BuildExtension 14 | }) 15 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/syncbn_gpu.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: syncbn-gpu 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/syncbn_gpu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/gpu/syncbn_gpu.so -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/syncbn_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include "common.h" 7 | #include "device_tensor.h" 8 | 9 | namespace { 10 | 11 | template 12 | struct GradOp { 13 | __device__ GradOp(Acctype m, const DeviceTensor3 i, const DeviceTensor3 g) 14 | : mean(m), input(i), gradOutput(g) {} 15 | __device__ __forceinline__ Float2 operator()(int batch, int plane, int n) { 16 | DType g = gradOutput[batch][plane][n]; 17 | DType c = ScalarConvert::to(input[batch][plane][n] - mean); 18 | return Float2(g, g * c); 19 | } 20 | const Acctype mean; 21 | const DeviceTensor3 input; 22 | const DeviceTensor3 gradOutput; 23 | }; 24 | 25 | template 26 | struct SumOp { 27 | __device__ SumOp(DeviceTensor i) : input(i){} 28 | __device__ __forceinline__ Float2 operator()(int batch, int plane, int n) { 29 | DType g = input[batch][plane][n]; 30 | return Float2(g, g * g); 31 | } 32 | DType mean; 33 | DeviceTensor input; 34 | }; 35 | 36 | // Sum across (batch, x/y/z) applying Op() pointwise 37 | template 38 | __device__ T reduce(Op op, DeviceTensor3 tensor, int plane) { 39 | T sum = (T)0; 40 | for (int batch = 0; batch < tensor.getSize(0); ++batch) { 41 | for (int x = threadIdx.x; x < tensor.getSize(2); x += blockDim.x) { 42 | sum += op(batch, plane, x); 43 | } 44 | } 45 | 46 | // sum over NumThreads within a warp 47 | sum = warpSum(sum); 48 | 49 | // 'transpose', and reduce within warp again 50 | __shared__ T shared[32]; 51 | __syncthreads(); 52 | if (threadIdx.x % WARP_SIZE == 0) { 53 | shared[threadIdx.x / WARP_SIZE] = sum; 54 | } 55 | if (threadIdx.x >= blockDim.x / WARP_SIZE && threadIdx.x < WARP_SIZE) { 56 | // zero out the other entries in shared 57 | shared[threadIdx.x] = (T)0; 58 | } 59 | __syncthreads(); 60 | if (threadIdx.x / WARP_SIZE == 0) { 61 | sum = warpSum(shared[threadIdx.x]); 62 | if (threadIdx.x == 0) { 63 | shared[0] = sum; 64 | } 65 | } 66 | __syncthreads(); 67 | 68 | // Everyone picks it up, should be broadcast into the whole gradInput 69 | return shared[0]; 70 | } 71 | 72 | template 73 | __global__ void BatchNorm_Forward_kernel ( 74 | DeviceTensor output, 75 | DeviceTensor input, 76 | DeviceTensor mean, 77 | DeviceTensor std, 78 | DeviceTensor gamma, 79 | DeviceTensor beta) { 80 | int c = blockIdx.x; 81 | /* main operation */ 82 | for (int b = 0; b < input.getSize(0); ++b) { 83 | for (int x = threadIdx.x; x < input.getSize(2); x += blockDim.x) { 84 | DType inp = input[b][c][x]; 85 | output[b][c][x] = gamma[c] * (inp - mean[c]) / 86 | std[c] + beta[c]; 87 | } 88 | } 89 | } 90 | 91 | template 92 | __global__ void BatchNorm_Backward_kernel ( 93 | DeviceTensor gradoutput, 94 | DeviceTensor input, 95 | DeviceTensor gradinput, 96 | DeviceTensor gradgamma, 97 | DeviceTensor gradbeta, 98 | DeviceTensor mean, 99 | DeviceTensor std, 100 | DeviceTensor gamma, 101 | DeviceTensor beta, 102 | DeviceTensor gradMean, 103 | DeviceTensor gradStd, 104 | bool train) { 105 | /* declarations of the variables */ 106 | /* Get the index and channels */ 107 | int c = blockIdx.x; 108 | /* main operation */ 109 | GradOp> g(mean[c], input, gradoutput); 110 | Float2 res = reduce, 111 | GradOp>, 112 | DeviceTensor>(g, gradoutput, c); 113 | DType gradOutputSum = res.v1; 114 | DType dotP = res.v2; 115 | DType invstd = DType(1.0) / std[c]; 116 | DType gradScale = invstd * gamma[c]; 117 | if (train && threadIdx.x == 0) { 118 | gradMean[c] = - gradOutputSum * gamma[c] * invstd; 119 | gradStd[c] = - dotP * gamma[c] * invstd * invstd; 120 | } 121 | if (gradinput.numElements() > 0) { 122 | for (int batch = 0; batch < gradoutput.getSize(0); ++batch) { 123 | for (int x = threadIdx.x; x < gradoutput.getSize(2); x += blockDim.x) { 124 | gradinput[batch][c][x] = gradoutput[batch][c][x] * gradScale; 125 | } 126 | } 127 | } 128 | if (gradgamma.numElements() > 0) { 129 | if (threadIdx.x == 0) { 130 | gradgamma[c] += dotP * invstd; 131 | } 132 | } 133 | if (gradbeta.numElements() > 0) { 134 | if (threadIdx.x == 0) { 135 | gradbeta[c] += gradOutputSum; 136 | } 137 | } 138 | } 139 | 140 | 141 | template 142 | __global__ void Sum_Square_Forward_kernel ( 143 | DeviceTensor input, 144 | DeviceTensor sum, 145 | DeviceTensor square) { 146 | int c = blockIdx.x; 147 | /* main operation */ 148 | SumOp g(input); 149 | Float2 res = reduce, 150 | SumOp, DeviceTensor>(g, input, c); 151 | DType xsum = res.v1; 152 | DType xsquare = res.v2; 153 | if (threadIdx.x == 0) { 154 | sum[c] = xsum; 155 | square[c] = xsquare; 156 | } 157 | } 158 | 159 | template 160 | __global__ void Sum_Square_Backward_kernel ( 161 | DeviceTensor gradInput, 162 | DeviceTensor input, 163 | DeviceTensor gradSum, 164 | DeviceTensor gradSquare) { 165 | int c = blockIdx.x; 166 | /* main operation */ 167 | for (int batch = 0; batch < gradInput.getSize(0); ++batch) { 168 | for (int x = threadIdx.x; x < gradInput.getSize(2); x += blockDim.x) 169 | { 170 | gradInput[batch][c][x] = gradSum[c] + 2 * gradSquare[c] * 171 | input[batch][c][x]; 172 | } 173 | } 174 | } 175 | 176 | } // namespcae 177 | 178 | at::Tensor BatchNorm_Forward_CUDA( 179 | const at::Tensor input_, 180 | const at::Tensor mean_, 181 | const at::Tensor std_, 182 | const at::Tensor gamma_, 183 | const at::Tensor beta_) { 184 | auto output_ = at::zeros_like(input_); 185 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 186 | dim3 blocks(input_.size(1)); 187 | dim3 threads(getNumThreads(input_.size(2))); 188 | AT_DISPATCH_FLOATING_TYPES(input_.type(), "BatchNorm_Forward_CUDA", ([&] { 189 | /* Device tensors */ 190 | DeviceTensor output = devicetensor(output_); 191 | DeviceTensor input = devicetensor(input_); 192 | DeviceTensor mean = devicetensor(mean_); 193 | DeviceTensor std = devicetensor(std_); 194 | DeviceTensor gamma = devicetensor(gamma_); 195 | DeviceTensor beta = devicetensor(beta_); 196 | /* kernel function */ 197 | BatchNorm_Forward_kernel<<>>( 198 | output, input, mean, std, gamma, beta); 199 | })); 200 | AT_ASSERT(cudaGetLastError() == cudaSuccess); 201 | return output_; 202 | } 203 | 204 | std::vector BatchNorm_Backward_CUDA( 205 | const at::Tensor gradoutput_, 206 | const at::Tensor input_, 207 | const at::Tensor mean_, 208 | const at::Tensor std_, 209 | const at::Tensor gamma_, 210 | const at::Tensor beta_, 211 | bool train) { 212 | /* outputs*/ 213 | at::Tensor gradinput_ = at::zeros_like(input_); 214 | at::Tensor gradgamma_ = at::zeros_like(gamma_); 215 | at::Tensor gradbeta_ = at::zeros_like(beta_); 216 | at::Tensor gradMean_ = at::zeros_like(mean_); 217 | at::Tensor gradStd_ = at::zeros_like(std_); 218 | /* cuda utils*/ 219 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 220 | dim3 blocks(input_.size(1)); 221 | dim3 threads(getNumThreads(input_.size(2))); 222 | AT_DISPATCH_FLOATING_TYPES(input_.type(), "BatchNorm_Backward_CUDA", ([&] { 223 | /* Device tensors */ 224 | DeviceTensor gradoutput = devicetensor(gradoutput_); 225 | DeviceTensor input = devicetensor(input_); 226 | DeviceTensor gradinput = devicetensor(gradinput_); 227 | DeviceTensor gradgamma = devicetensor(gradgamma_); 228 | DeviceTensor gradbeta = devicetensor(gradbeta_); 229 | DeviceTensor mean = devicetensor(mean_); 230 | DeviceTensor std = devicetensor(std_); 231 | DeviceTensor gamma = devicetensor(gamma_); 232 | DeviceTensor beta = devicetensor(beta_); 233 | DeviceTensor gradMean = devicetensor(gradMean_); 234 | DeviceTensor gradStd = devicetensor(gradStd_); 235 | /* kernel function */ 236 | BatchNorm_Backward_kernel 237 | <<>>( 238 | gradoutput, input, gradinput, gradgamma, gradbeta, mean, std, 239 | gamma, beta, gradMean, gradStd, train); 240 | })); 241 | AT_ASSERT(cudaGetLastError() == cudaSuccess); 242 | return {gradinput_, gradMean_, gradStd_, gradgamma_, gradbeta_}; 243 | } 244 | 245 | std::vector Sum_Square_Forward_CUDA( 246 | const at::Tensor input_) { 247 | /* outputs */ 248 | at::Tensor sum_ = torch::zeros({input_.size(1)}, input_.options()); 249 | at::Tensor square_ = torch::zeros({input_.size(1)}, input_.options()); 250 | /* cuda utils*/ 251 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 252 | dim3 blocks(input_.size(1)); 253 | dim3 threads(getNumThreads(input_.size(2))); 254 | AT_DISPATCH_FLOATING_TYPES(input_.type(), "SumSquare_forward_CUDA", ([&] { 255 | /* Device tensors */ 256 | DeviceTensor input = devicetensor(input_); 257 | DeviceTensor sum = devicetensor(sum_); 258 | DeviceTensor square = devicetensor(square_); 259 | /* kernel function */ 260 | Sum_Square_Forward_kernel 261 | <<>>(input, sum, square); 262 | })); 263 | AT_ASSERT(cudaGetLastError() == cudaSuccess); 264 | return {sum_, square_}; 265 | } 266 | 267 | at::Tensor Sum_Square_Backward_CUDA( 268 | const at::Tensor input_, 269 | const at::Tensor gradSum_, 270 | const at::Tensor gradSquare_) { 271 | /* outputs */ 272 | at::Tensor gradInput_ = at::zeros_like(input_); 273 | /* cuda utils*/ 274 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 275 | dim3 blocks(input_.size(1)); 276 | dim3 threads(getNumThreads(input_.size(2))); 277 | AT_DISPATCH_FLOATING_TYPES(input_.type(), "SumSquare_Backward_CUDA", ([&] { 278 | /* Device tensors */ 279 | DeviceTensor gradInput = devicetensor(gradInput_); 280 | DeviceTensor input = devicetensor(input_); 281 | DeviceTensor gradSum = devicetensor(gradSum_); 282 | DeviceTensor gradSquare =devicetensor(gradSquare_); 283 | /* kernel function */ 284 | Sum_Square_Backward_kernel 285 | <<>>(gradInput, input, gradSum, gradSquare); 286 | })); 287 | AT_ASSERT(cudaGetLastError() == cudaSuccess); 288 | return gradInput_; 289 | } 290 | -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/src/gpu/syncbn_kernel.cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/seg_opr/sync_bn/src/gpu/syncbn_kernel.cuda.o -------------------------------------------------------------------------------- /furnace/seg_opr/sync_bn/syncbn.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/10/3 下午1:45 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : syncbn.py.py 7 | 8 | """Synchronized Cross-GPU Batch Normalization Module""" 9 | import collections 10 | import threading 11 | 12 | import torch 13 | from torch.nn.modules.batchnorm import _BatchNorm 14 | from torch.nn.functional import batch_norm 15 | from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast 16 | 17 | from .functions import * 18 | from .comm import SyncMaster 19 | from .parallel import allreduce 20 | 21 | 22 | __all__ = ['BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d'] 23 | 24 | class _SyncBatchNorm(_BatchNorm): 25 | def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True): 26 | super(_SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine) 27 | 28 | self._sync_master = SyncMaster(self._data_parallel_master) 29 | self._parallel_id = None 30 | self._slave_pipe = None 31 | 32 | def forward(self, input): 33 | if not self.training: 34 | return batch_norm( 35 | input, self.running_mean, self.running_var, self.weight, self.bias, 36 | self.training, self.momentum, self.eps) 37 | 38 | # Resize the input to (B, C, -1). 39 | input_shape = input.size() 40 | input = input.view(input_shape[0], self.num_features, -1) 41 | 42 | # sum(x) and sum(x^2) 43 | N = input.size(0) * input.size(2) 44 | xsum, xsqsum = sum_square(input) 45 | 46 | # all-reduce for global sum(x) and sum(x^2) 47 | if self._parallel_id == 0: 48 | mean, inv_std = self._sync_master.run_master(_ChildMessage(xsum, xsqsum, N)) 49 | else: 50 | mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(xsum, xsqsum, N)) 51 | # forward 52 | return batchnormtrain(input, mean, 1.0/inv_std, self.weight, self.bias).view(input_shape) 53 | 54 | def __data_parallel_replicate__(self, ctx, copy_id): 55 | self._parallel_id = copy_id 56 | 57 | # parallel_id == 0 means master device. 58 | if self._parallel_id == 0: 59 | ctx.sync_master = self._sync_master 60 | else: 61 | self._slave_pipe = ctx.sync_master.register_slave(copy_id) 62 | 63 | def _data_parallel_master(self, intermediates): 64 | """Reduce the sum and square-sum, compute the statistics, and broadcast it.""" 65 | 66 | # Always using same "device order" makes the ReduceAdd operation faster. 67 | # Thanks to:: Tete Xiao (http://tetexiao.com/) 68 | intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device()) 69 | 70 | to_reduce = [i[1][:2] for i in intermediates] 71 | to_reduce = [j for i in to_reduce for j in i] # flatten 72 | target_gpus = [i[1].sum.get_device() for i in intermediates] 73 | 74 | sum_size = sum([i[1].sum_size for i in intermediates]) 75 | sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce) 76 | mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size) 77 | 78 | broadcasted = Broadcast.apply(target_gpus, mean, inv_std) 79 | 80 | outputs = [] 81 | for i, rec in enumerate(intermediates): 82 | outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2]))) 83 | 84 | return outputs 85 | 86 | def _compute_mean_std(self, sum_, ssum, size): 87 | """Compute the mean and standard-deviation with sum and square-sum. This method 88 | also maintains the moving average on the master device.""" 89 | assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.' 90 | mean = sum_ / size 91 | sumvar = ssum - sum_ * mean 92 | unbias_var = sumvar / (size - 1) 93 | bias_var = sumvar / size 94 | 95 | self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data 96 | self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data 97 | 98 | return mean, (bias_var + self.eps) ** -0.5 99 | 100 | 101 | # API adapted from https://github.com/vacancy/Synchronized-BatchNorm-PyTorch 102 | _ChildMessage = collections.namedtuple('Message', ['sum', 'ssum', 'sum_size']) 103 | _MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std']) 104 | 105 | 106 | class BatchNorm1d(_SyncBatchNorm): 107 | r"""Please see the docs in :class:`encoding.nn.BatchNorm2d`""" 108 | def _check_input_dim(self, input): 109 | if input.dim() != 2 and input.dim() != 3: 110 | raise ValueError('expected 2D or 3D input (got {}D input)' 111 | .format(input.dim())) 112 | super(BatchNorm2d, self)._check_input_dim(input) 113 | 114 | 115 | class BatchNorm2d(_SyncBatchNorm): 116 | r"""Cross-GPU Synchronized Batch normalization (SyncBN) 117 | 118 | Standard BN [1]_ implementation only normalize the data within each device (GPU). 119 | SyncBN normalizes the input within the whole mini-batch. 120 | We follow the sync-onece implmentation described in the paper [2]_ . 121 | Please see the design idea in the `notes <./notes/syncbn.html>`_. 122 | 123 | .. note:: 124 | We adapt the awesome python API from another `PyTorch SyncBN Implementation 125 | `_ and provide 126 | efficient CUDA backend. 127 | 128 | .. math:: 129 | 130 | y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta 131 | 132 | The mean and standard-deviation are calculated per-channel over 133 | the mini-batches and gamma and beta are learnable parameter vectors 134 | of size C (where C is the input size). 135 | 136 | During training, this layer keeps a running estimate of its computed mean 137 | and variance. The running sum is kept with a default momentum of 0.1. 138 | 139 | During evaluation, this running mean/variance is used for normalization. 140 | 141 | Because the BatchNorm is done over the `C` dimension, computing statistics 142 | on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm 143 | 144 | Args: 145 | num_features: num_features from an expected input of 146 | size batch_size x num_features x height x width 147 | eps: a value added to the denominator for numerical stability. 148 | Default: 1e-5 149 | momentum: the value used for the running_mean and running_var 150 | computation. Default: 0.1 151 | affine: a boolean value that when set to ``True``, gives the layer learnable 152 | affine parameters. Default: ``True`` 153 | 154 | Shape: 155 | - Input: :math:`(N, C, H, W)` 156 | - Output: :math:`(N, C, H, W)` (same shape as input) 157 | 158 | Reference: 159 | .. [1] Ioffe, Sergey, and Christian Szegedy. "Batch normalization: Accelerating deep network training by reducing internal covariate shift." *ICML 2015* 160 | .. [2] Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi, and Amit Agrawal. "Context Encoding for Semantic Segmentation." *CVPR 2018* 161 | 162 | Examples: 163 | >>> m = BatchNorm2d(100) 164 | >>> net = torch.nn.DataParallel(m) 165 | >>> encoding.parallel.patch_replication_callback(net) 166 | >>> output = net(input) 167 | """ 168 | def _check_input_dim(self, input): 169 | if input.dim() != 4: 170 | raise ValueError('expected 4D input (got {}D input)' 171 | .format(input.dim())) 172 | super(BatchNorm2d, self)._check_input_dim(input) 173 | 174 | 175 | class BatchNorm3d(_SyncBatchNorm): 176 | r"""Please see the docs in :class:`encoding.nn.BatchNorm2d`""" 177 | def _check_input_dim(self, input): 178 | if input.dim() != 5: 179 | raise ValueError('expected 5D input (got {}D input)' 180 | .format(input.dim())) 181 | super(BatchNorm3d, self)._check_input_dim(input) 182 | 183 | 184 | class SharedTensor(object): 185 | """Shared Tensor for cross GPU all reduce operation""" 186 | def __init__(self, nGPUs): 187 | self.mutex = threading.Lock() 188 | self.all_tasks_done = threading.Condition(self.mutex) 189 | self.nGPUs = nGPUs 190 | self._clear() 191 | 192 | def _clear(self): 193 | self.N = 0 194 | self.dict = {} 195 | self.push_tasks = self.nGPUs 196 | self.reduce_tasks = self.nGPUs 197 | 198 | def push(self, *inputs): 199 | # push from device 200 | with self.mutex: 201 | if self.push_tasks == 0: 202 | self._clear() 203 | self.N += inputs[0] 204 | igpu = inputs[1] 205 | self.dict[igpu] = inputs[2:] 206 | #idx = self.nGPUs - self.push_tasks 207 | self.push_tasks -= 1 208 | with self.all_tasks_done: 209 | if self.push_tasks == 0: 210 | self.all_tasks_done.notify_all() 211 | while self.push_tasks: 212 | self.all_tasks_done.wait() 213 | 214 | def pull(self, igpu): 215 | # pull from device 216 | with self.mutex: 217 | if igpu == 0: 218 | assert(len(self.dict) == self.nGPUs) 219 | # flatten the tensors 220 | self.list = [t for i in range(len(self.dict)) for t in self.dict[i]] 221 | self.outlist = allreduce(2, *self.list) 222 | self.reduce_tasks -= 1 223 | else: 224 | self.reduce_tasks -= 1 225 | with self.all_tasks_done: 226 | if self.reduce_tasks == 0: 227 | self.all_tasks_done.notify_all() 228 | while self.reduce_tasks: 229 | self.all_tasks_done.wait() 230 | # all reduce done 231 | return self.N, self.outlist[2*igpu], self.outlist[2*igpu+1] 232 | 233 | def __len__(self): 234 | return self.nGPUs 235 | 236 | def __repr__(self): 237 | return ('SharedTensor') 238 | -------------------------------------------------------------------------------- /furnace/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__init__.py -------------------------------------------------------------------------------- /furnace/utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/img_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/img_utils.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/img_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/img_utils.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/img_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/img_utils.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/init_func.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/init_func.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/init_func.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/init_func.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/init_func.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/init_func.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/pyt_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/pyt_utils.cpython-36.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/pyt_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/pyt_utils.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/pyt_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/pyt_utils.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/visualize.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/visualize.cpython-37.pyc -------------------------------------------------------------------------------- /furnace/utils/__pycache__/visualize.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjcaimeow/SISNet/9822aa29cae7690fbde6bd260c67346919543768/furnace/utils/__pycache__/visualize.cpython-38.pyc -------------------------------------------------------------------------------- /furnace/utils/img_utils.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import numbers 4 | import random 5 | import collections 6 | 7 | 8 | def get_2dshape(shape, *, zero=True): 9 | if not isinstance(shape, collections.Iterable): 10 | shape = int(shape) 11 | shape = (shape, shape) 12 | else: 13 | h, w = map(int, shape) 14 | shape = (h, w) 15 | if zero: 16 | minv = 0 17 | else: 18 | minv = 1 19 | 20 | assert min(shape) >= minv, 'invalid shape: {}'.format(shape) 21 | return shape 22 | 23 | 24 | def random_crop_pad_to_shape(img, crop_pos, crop_size, pad_label_value): 25 | h, w = img.shape[:2] 26 | start_crop_h, start_crop_w = crop_pos 27 | assert ((start_crop_h < h) and (start_crop_h >= 0)) 28 | assert ((start_crop_w < w) and (start_crop_w >= 0)) 29 | 30 | crop_size = get_2dshape(crop_size) 31 | crop_h, crop_w = crop_size 32 | 33 | img_crop = img[start_crop_h:start_crop_h + crop_h, 34 | start_crop_w:start_crop_w + crop_w, ...] 35 | 36 | img_, margin = pad_image_to_shape(img_crop, crop_size, cv2.BORDER_CONSTANT, 37 | pad_label_value) 38 | 39 | return img_, margin 40 | 41 | 42 | def generate_random_crop_pos(ori_size, crop_size): 43 | ori_size = get_2dshape(ori_size) 44 | h, w = ori_size 45 | 46 | crop_size = get_2dshape(crop_size) 47 | crop_h, crop_w = crop_size 48 | 49 | pos_h, pos_w = 0, 0 50 | 51 | if h > crop_h: 52 | pos_h = random.randint(0, h - crop_h + 1) 53 | 54 | if w > crop_w: 55 | pos_w = random.randint(0, w - crop_w + 1) 56 | 57 | return pos_h, pos_w 58 | 59 | 60 | def pad_image_to_shape(img, shape, border_mode, value): 61 | margin = np.zeros(4, np.uint32) 62 | shape = get_2dshape(shape) 63 | pad_height = shape[0] - img.shape[0] if shape[0] - img.shape[0] > 0 else 0 64 | pad_width = shape[1] - img.shape[1] if shape[1] - img.shape[1] > 0 else 0 65 | 66 | margin[0] = pad_height // 2 67 | margin[1] = pad_height // 2 + pad_height % 2 68 | margin[2] = pad_width // 2 69 | margin[3] = pad_width // 2 + pad_width % 2 70 | 71 | img = cv2.copyMakeBorder(img, margin[0], margin[1], margin[2], margin[3], 72 | border_mode, value=value) 73 | 74 | return img, margin 75 | 76 | 77 | def pad_image_size_to_multiples_of(img, multiple, pad_value): 78 | h, w = img.shape[:2] 79 | d = multiple 80 | 81 | def canonicalize(s): 82 | v = s // d 83 | return (v + (v * d != s)) * d 84 | 85 | th, tw = map(canonicalize, (h, w)) 86 | 87 | return pad_image_to_shape(img, (th, tw), cv2.BORDER_CONSTANT, pad_value) 88 | 89 | 90 | def resize_ensure_shortest_edge(img, edge_length, 91 | interpolation_mode=cv2.INTER_LINEAR): 92 | assert isinstance(edge_length, int) and edge_length > 0, edge_length 93 | h, w = img.shape[:2] 94 | if h < w: 95 | ratio = float(edge_length) / h 96 | th, tw = edge_length, max(1, int(ratio * w)) 97 | else: 98 | ratio = float(edge_length) / w 99 | th, tw = max(1, int(ratio * h)), edge_length 100 | img = cv2.resize(img, (tw, th), interpolation_mode) 101 | 102 | return img 103 | 104 | 105 | def random_scale(img, gt, scales): 106 | scale = random.choice(scales) 107 | sh = int(img.shape[0] * scale) 108 | sw = int(img.shape[1] * scale) 109 | img = cv2.resize(img, (sw, sh), interpolation=cv2.INTER_LINEAR) 110 | gt = cv2.resize(gt, (sw, sh), interpolation=cv2.INTER_NEAREST) 111 | 112 | return img, gt, scale 113 | 114 | 115 | def random_scale_with_length(img, gt, length): 116 | size = random.choice(length) 117 | sh = size 118 | sw = size 119 | img = cv2.resize(img, (sw, sh), interpolation=cv2.INTER_LINEAR) 120 | gt = cv2.resize(gt, (sw, sh), interpolation=cv2.INTER_NEAREST) 121 | 122 | return img, gt, size 123 | 124 | 125 | def random_mirror(img, gt): 126 | if random.random() >= 0.5: 127 | img = cv2.flip(img, 1) 128 | gt = cv2.flip(gt, 1) 129 | 130 | return img, gt, 131 | 132 | 133 | def random_rotation(img, gt): 134 | angle = random.random() * 20 - 10 135 | h, w = img.shape[:2] 136 | rotation_matrix = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1) 137 | img = cv2.warpAffine(img, rotation_matrix, (w, h), flags=cv2.INTER_LINEAR) 138 | gt = cv2.warpAffine(gt, rotation_matrix, (w, h), flags=cv2.INTER_NEAREST) 139 | 140 | return img, gt 141 | 142 | 143 | def random_gaussian_blur(img): 144 | gauss_size = random.choice([1, 3, 5, 7]) 145 | if gauss_size > 1: 146 | # do the gaussian blur 147 | img = cv2.GaussianBlur(img, (gauss_size, gauss_size), 0) 148 | 149 | return img 150 | 151 | 152 | def center_crop(img, shape): 153 | h, w = shape[0], shape[1] 154 | y = (img.shape[0] - h) // 2 155 | x = (img.shape[1] - w) // 2 156 | return img[y:y + h, x:x + w] 157 | 158 | 159 | def random_crop(img, gt, size): 160 | if isinstance(size, numbers.Number): 161 | size = (int(size), int(size)) 162 | else: 163 | size = size 164 | 165 | h, w = img.shape[:2] 166 | crop_h, crop_w = size[0], size[1] 167 | 168 | if h > crop_h: 169 | x = random.randint(0, h - crop_h + 1) 170 | img = img[x:x + crop_h, :, :] 171 | gt = gt[x:x + crop_h, :] 172 | 173 | if w > crop_w: 174 | x = random.randint(0, w - crop_w + 1) 175 | img = img[:, x:x + crop_w, :] 176 | gt = gt[:, x:x + crop_w] 177 | 178 | return img, gt 179 | 180 | 181 | def normalize(img, mean, std): 182 | # pytorch pretrained model need the input range: 0-1 183 | img = img.astype(np.float32) / 255.0 184 | img = img - mean 185 | img = img / std 186 | 187 | return img 188 | -------------------------------------------------------------------------------- /furnace/utils/init_func.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/9/28 下午12:13 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : init_func.py.py 7 | import torch 8 | import torch.nn as nn 9 | 10 | 11 | def __init_weight(feature, conv_init, norm_layer, bn_eps, bn_momentum, 12 | **kwargs): 13 | for name, m in feature.named_modules(): 14 | if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 15 | conv_init(m.weight, **kwargs) 16 | elif isinstance(m, norm_layer): 17 | m.eps = bn_eps 18 | m.momentum = bn_momentum 19 | nn.init.constant_(m.weight, 1) 20 | nn.init.constant_(m.bias, 0) 21 | 22 | 23 | def init_weight(module_list, conv_init, norm_layer, bn_eps, bn_momentum, 24 | **kwargs): 25 | if isinstance(module_list, list): 26 | for feature in module_list: 27 | __init_weight(feature, conv_init, norm_layer, bn_eps, bn_momentum, 28 | **kwargs) 29 | else: 30 | __init_weight(module_list, conv_init, norm_layer, bn_eps, bn_momentum, 31 | **kwargs) 32 | 33 | 34 | def group_weight(weight_group, module, norm_layer, lr): 35 | group_decay = [] 36 | group_no_decay = [] 37 | for m in module.modules(): 38 | if isinstance(m, nn.Linear): 39 | group_decay.append(m.weight) 40 | if m.bias is not None: 41 | group_no_decay.append(m.bias) 42 | elif isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.ConvTranspose2d, nn.ConvTranspose3d)): 43 | group_decay.append(m.weight) 44 | if m.bias is not None: 45 | group_no_decay.append(m.bias) 46 | elif isinstance(m, norm_layer) or isinstance(m, nn.BatchNorm1d) or isinstance(m, nn.BatchNorm2d) \ 47 | or isinstance(m, nn.BatchNorm3d) or isinstance(m, nn.GroupNorm): 48 | if m.weight is not None: 49 | group_no_decay.append(m.weight) 50 | if m.bias is not None: 51 | group_no_decay.append(m.bias) 52 | elif isinstance(m, nn.Parameter): 53 | group_decay.append(m) 54 | # else: 55 | # print(m, norm_layer) 56 | # print(module.modules) 57 | # print( len(list(module.parameters())) , 'HHHHHHHHHHHHHHHHH', len(group_decay) + len( 58 | # group_no_decay)) 59 | assert len(list(module.parameters())) == len(group_decay) + len( 60 | group_no_decay) 61 | weight_group.append(dict(params=group_decay, lr=lr)) 62 | weight_group.append(dict(params=group_no_decay, weight_decay=.0, lr=lr)) 63 | return weight_group 64 | -------------------------------------------------------------------------------- /furnace/utils/pyt_utils.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | import os 3 | import sys 4 | import time 5 | import random 6 | import argparse 7 | import logging 8 | from collections import OrderedDict, defaultdict 9 | 10 | import torch 11 | import torch.utils.model_zoo as model_zoo 12 | 13 | class LogFormatter(logging.Formatter): 14 | log_fout = None 15 | date_full = '[%(asctime)s %(lineno)d@%(filename)s:%(name)s] ' 16 | date = '%(asctime)s ' 17 | msg = '%(message)s' 18 | 19 | def format(self, record): 20 | if record.levelno == logging.DEBUG: 21 | mcl, mtxt = self._color_dbg, 'DBG' 22 | elif record.levelno == logging.WARNING: 23 | mcl, mtxt = self._color_warn, 'WRN' 24 | elif record.levelno == logging.ERROR: 25 | mcl, mtxt = self._color_err, 'ERR' 26 | else: 27 | mcl, mtxt = self._color_normal, '' 28 | 29 | if mtxt: 30 | mtxt += ' ' 31 | 32 | if self.log_fout: 33 | self.__set_fmt(self.date_full + mtxt + self.msg) 34 | formatted = super(LogFormatter, self).format(record) 35 | # self.log_fout.write(formatted) 36 | # self.log_fout.write('\n') 37 | # self.log_fout.flush() 38 | return formatted 39 | 40 | self.__set_fmt(self._color_date(self.date) + mcl(mtxt + self.msg)) 41 | formatted = super(LogFormatter, self).format(record) 42 | 43 | return formatted 44 | 45 | if sys.version_info.major < 3: 46 | def __set_fmt(self, fmt): 47 | self._fmt = fmt 48 | else: 49 | def __set_fmt(self, fmt): 50 | self._style._fmt = fmt 51 | 52 | @staticmethod 53 | def _color_dbg(msg): 54 | return '\x1b[36m{}\x1b[0m'.format(msg) 55 | 56 | @staticmethod 57 | def _color_warn(msg): 58 | return '\x1b[1;31m{}\x1b[0m'.format(msg) 59 | 60 | @staticmethod 61 | def _color_err(msg): 62 | return '\x1b[1;4;31m{}\x1b[0m'.format(msg) 63 | 64 | @staticmethod 65 | def _color_omitted(msg): 66 | return '\x1b[35m{}\x1b[0m'.format(msg) 67 | 68 | @staticmethod 69 | def _color_normal(msg): 70 | return msg 71 | 72 | @staticmethod 73 | def _color_date(msg): 74 | return '\x1b[32m{}\x1b[0m'.format(msg) 75 | 76 | _default_level_name = os.getenv('ENGINE_LOGGING_LEVEL', 'INFO') 77 | _default_level = logging.getLevelName(_default_level_name.upper()) 78 | 79 | def get_logger(log_dir=None, log_file=None, formatter=LogFormatter): 80 | logger = logging.getLogger() 81 | logger.setLevel(_default_level) 82 | del logger.handlers[:] 83 | 84 | if log_dir and log_file: 85 | ensure_dir(log_dir) 86 | LogFormatter.log_fout = True 87 | file_handler = logging.FileHandler(log_file, mode='a') 88 | file_handler.setLevel(logging.INFO) 89 | file_handler.setFormatter(formatter) 90 | logger.addHandler(file_handler) 91 | 92 | stream_handler = logging.StreamHandler() 93 | stream_handler.setFormatter(formatter(datefmt='%d %H:%M:%S')) 94 | stream_handler.setLevel(0) 95 | logger.addHandler(stream_handler) 96 | return logger 97 | 98 | logger = get_logger() 99 | 100 | model_urls = { 101 | 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 102 | 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 103 | 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 104 | 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 105 | 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 106 | } 107 | 108 | 109 | def load_model(model, model_file, is_restore=False): 110 | t_start = time.time() 111 | 112 | if model_file is None: 113 | return model 114 | 115 | if isinstance(model_file, str): 116 | state_dict = torch.load(model_file) 117 | if 'model' in state_dict.keys(): 118 | state_dict = state_dict['model'] 119 | else: 120 | state_dict = model_file 121 | t_ioend = time.time() 122 | 123 | if is_restore: 124 | new_state_dict = OrderedDict() 125 | for k, v in state_dict.items(): 126 | name = 'module.' + k 127 | new_state_dict[name] = v 128 | state_dict = new_state_dict 129 | 130 | 131 | model.load_state_dict(state_dict, strict=False) 132 | ckpt_keys = set(state_dict.keys()) 133 | own_keys = set(model.state_dict().keys()) 134 | missing_keys = own_keys - ckpt_keys 135 | unexpected_keys = ckpt_keys - own_keys 136 | 137 | del state_dict 138 | t_end = time.time() 139 | logger.info( 140 | "Load model, Time usage:\n\tIO: {}, initialize parameters: {}".format( 141 | t_ioend - t_start, t_end - t_ioend)) 142 | 143 | return model 144 | 145 | def load_dualpath_model(model, model_file, is_restore=False): 146 | # load raw state_dict 147 | t_start = time.time() 148 | if isinstance(model_file, str): 149 | raw_state_dict = torch.load(model_file) 150 | 151 | 152 | if 'model' in raw_state_dict.keys(): 153 | raw_state_dict = raw_state_dict['model'] 154 | else: 155 | raw_state_dict = model_file 156 | # copy to hha backbone 157 | state_dict = {} 158 | for k, v in raw_state_dict.items(): 159 | state_dict[k.replace('.bn.', '.')] = v 160 | if k.find('conv1') >= 0: 161 | state_dict[k] = v 162 | state_dict[k.replace('conv1', 'hha_conv1')] = v 163 | if k.find('conv2') >= 0: 164 | state_dict[k] = v 165 | state_dict[k.replace('conv2', 'hha_conv2')] = v 166 | if k.find('conv3') >= 0: 167 | state_dict[k] = v 168 | state_dict[k.replace('conv3', 'hha_conv3')] = v 169 | if k.find('bn1') >= 0: 170 | state_dict[k] = v 171 | state_dict[k.replace('bn1', 'hha_bn1')] = v 172 | if k.find('bn2') >= 0: 173 | state_dict[k] = v 174 | state_dict[k.replace('bn2', 'hha_bn2')] = v 175 | if k.find('bn3') >= 0: 176 | state_dict[k] = v 177 | state_dict[k.replace('bn3', 'hha_bn3')] = v 178 | if k.find('downsample') >= 0: 179 | state_dict[k] = v 180 | state_dict[k.replace('downsample', 'hha_downsample')] = v 181 | t_ioend = time.time() 182 | 183 | if is_restore: 184 | new_state_dict = OrderedDict() 185 | for k, v in state_dict.items(): 186 | name = 'module.' + k 187 | new_state_dict[name] = v 188 | state_dict = new_state_dict 189 | 190 | model.load_state_dict(state_dict, strict=False) 191 | ckpt_keys = set(state_dict.keys()) 192 | own_keys = set(model.state_dict().keys()) 193 | missing_keys = own_keys - ckpt_keys 194 | unexpected_keys = ckpt_keys - own_keys 195 | 196 | # if len(missing_keys) > 0: 197 | # logger.warning('Missing key(s) in state_dict: {}'.format( 198 | # ', '.join('{}'.format(k) for k in missing_keys))) 199 | # 200 | # if len(unexpected_keys) > 0: 201 | # logger.warning('Unexpected key(s) in state_dict: {}'.format( 202 | # ', '.join('{}'.format(k) for k in unexpected_keys))) 203 | 204 | del state_dict 205 | t_end = time.time() 206 | logger.info( 207 | "Load model, Time usage:\n\tIO: {}, initialize parameters: {}".format( 208 | t_ioend - t_start, t_end - t_ioend)) 209 | 210 | return model 211 | 212 | def parse_devices(input_devices): 213 | if input_devices.endswith('*'): 214 | devices = list(range(torch.cuda.device_count())) 215 | return devices 216 | 217 | devices = [] 218 | for d in input_devices.split(','): 219 | if '-' in d: 220 | start_device, end_device = d.split('-')[0], d.split('-')[1] 221 | assert start_device != '' 222 | assert end_device != '' 223 | start_device, end_device = int(start_device), int(end_device) 224 | assert start_device < end_device 225 | assert end_device < torch.cuda.device_count() 226 | for sd in range(start_device, end_device + 1): 227 | devices.append(sd) 228 | else: 229 | device = int(d) 230 | assert device < torch.cuda.device_count() 231 | devices.append(device) 232 | 233 | logger.info('using devices {}'.format( 234 | ', '.join([str(d) for d in devices]))) 235 | 236 | return devices 237 | 238 | 239 | def extant_file(x): 240 | """ 241 | 'Type' for argparse - checks that file exists but does not open. 242 | """ 243 | if not os.path.exists(x): 244 | # Argparse uses the ArgumentTypeError to give a rejection message like: 245 | # error: argument input: x does not exist 246 | raise argparse.ArgumentTypeError("{0} does not exist".format(x)) 247 | return x 248 | 249 | 250 | def link_file(src, target): 251 | if os.path.isdir(target) or os.path.isfile(target): 252 | os.system('rm -rf {}'.format(target)) 253 | os.system('ln -s {} {}'.format(src, target)) 254 | 255 | 256 | def ensure_dir(path): 257 | if not os.path.isdir(path): 258 | try: 259 | sleeptime = random.randint(0, 3) 260 | time.sleep(sleeptime) 261 | os.makedirs(path) 262 | except: 263 | print('conflict !!!') 264 | 265 | 266 | def _dbg_interactive(var, value): 267 | from IPython import embed 268 | embed() 269 | 270 | # def load_model(model, model_file): 271 | # if isinstance(model_file, str): 272 | # print('Load Model: ' + model_file) 273 | # state_dict = torch.load(model_file) 274 | # else: 275 | # state_dict = model_file 276 | # 277 | # from collections import OrderedDict 278 | # new_state_dict = OrderedDict() 279 | # for k, v in state_dict.items(): 280 | # name = k 281 | # if k.split('.')[0] == 'module': 282 | # name = k[7:] 283 | # new_state_dict[name] = v 284 | # model.load_state_dict(new_state_dict, strict=False) 285 | # 286 | # return model 287 | # 288 | # 289 | # def parse_devices(input_devices): 290 | # if input_devices.endswith('*'): 291 | # devices = list(range(torch.cuda.device_count())) 292 | # return devices 293 | # 294 | # devices = [] 295 | # for d in input_devices.split(','): 296 | # if '-' in d: 297 | # start_device, end_device = d.split('-')[0], d.split('-')[1] 298 | # assert start_device != '' 299 | # assert end_device != '' 300 | # start_device, end_device = int(start_device), int(end_device) 301 | # assert start_device < end_device 302 | # assert end_device < torch.cuda.device_count() 303 | # for sd in range(start_device, end_device + 1): 304 | # devices.append(sd) 305 | # else: 306 | # device = int(d) 307 | # assert device < torch.cuda.device_count() 308 | # devices.append(device) 309 | # 310 | # return devices 311 | 312 | # 313 | # def inspect(var): 314 | # return CallbackInjector(var, _dbg_interactive) 315 | -------------------------------------------------------------------------------- /furnace/utils/visualize.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import scipy.io as sio 4 | 5 | def set_img_color(colors, background, img, pred, gt, show255=False): 6 | for i in range(0, len(colors)): 7 | if i != background: 8 | img[np.where(pred == i)] = colors[i] 9 | if show255: 10 | img[np.where(gt==background)] = 255 11 | return img 12 | 13 | def show_prediction(colors, background, img, pred, gt): 14 | im = np.array(img, np.uint8) 15 | set_img_color(colors, background, im, pred, gt) 16 | final = np.array(im) 17 | return final 18 | 19 | def show_img(colors, background, img, clean, gt, *pds): 20 | im1 = np.array(img, np.uint8) 21 | #set_img_color(colors, background, im1, clean, gt) 22 | final = np.array(im1) 23 | # the pivot black bar 24 | pivot = np.zeros((im1.shape[0], 15, 3), dtype=np.uint8) 25 | for pd in pds: 26 | im = np.array(img, np.uint8) 27 | # pd[np.where(gt == 255)] = 255 28 | set_img_color(colors, background, im, pd, gt) 29 | final = np.column_stack((final, pivot)) 30 | final = np.column_stack((final, im)) 31 | 32 | im = np.array(img, np.uint8) 33 | set_img_color(colors, background, im, gt, True) 34 | final = np.column_stack((final, pivot)) 35 | final = np.column_stack((final, im)) 36 | return final 37 | 38 | def get_colors(class_num): 39 | colors = [] 40 | for i in range(class_num): 41 | colors.append((np.random.random((1,3)) * 255).tolist()[0]) 42 | 43 | return colors 44 | 45 | def get_ade_colors(): 46 | colors = sio.loadmat('./color150.mat')['colors'] 47 | colors = colors[:,::-1,] 48 | colors = np.array(colors).astype(int).tolist() 49 | colors.insert(0,[0,0,0]) 50 | 51 | return colors 52 | 53 | def print_iou(iu, mean_pixel_acc, class_names=None, show_no_back=False, no_print=False): 54 | n = iu.size 55 | lines = [] 56 | for i in range(n): 57 | if class_names is None: 58 | cls = 'Class %d:' % (i+1) 59 | else: 60 | cls = '%d %s' % (i+1, class_names[i]) 61 | lines.append('%-8s\t%.3f%%' % (cls, iu[i] * 100)) 62 | mean_IU = np.nanmean(iu) 63 | mean_IU_no_back = np.nanmean(iu[1:]) 64 | if show_no_back: 65 | lines.append('---------------------------- %-8s\t%.3f%%\t%-8s\t%.3f%%\t%-8s\t%.3f%%' % ('mean_IU', mean_IU * 100, 'mean_IU_no_back', mean_IU_no_back*100, 66 | 'mean_pixel_ACC',mean_pixel_acc*100)) 67 | else: 68 | print(mean_pixel_acc) 69 | lines.append('---------------------------- %-8s\t%.3f%%\t%-8s\t%.3f%%' % ('mean_IU', mean_IU * 100,'mean_pixel_ACC',mean_pixel_acc*100)) 70 | line = "\n".join(lines) 71 | if not no_print: 72 | print(line) 73 | return line 74 | 75 | 76 | --------------------------------------------------------------------------------