├── .gitignore ├── LICENSE.md ├── README.md ├── bin └── dgcnn.py ├── dgcnn ├── __init__.py ├── flags.py ├── iotool.py ├── main_funcs.py ├── model.py ├── ops.py └── trainval.py └── scripts ├── lsf ├── continue_dgcnn_vnp.sh ├── inference_dgcnn.sh └── train_dgcnn.sh └── plot_csv.py /.gitignore: -------------------------------------------------------------------------------- 1 | arxiv 2 | *~ 3 | *.csv 4 | *.dat 5 | *.root 6 | *.pyc 7 | log 8 | *.txt 9 | weights -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Kazuhiro Terao 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # dgcnn 2 | 3 | This is an implementation of 3D point cloud semantic segmentation for [Dynamic Graph Convolutional Neural Network](https://arxiv.org/abs/1801.07829). The number of edge convolution layers, fully connected layers, and number of filters per each layer are all configurable. The implementation includes a few variations such as residual unit (edge convolution with identity mapping), with or without fully connected layers, etc.. Experimental results on DeepLearnPhysics open data set will be made available. 4 | 5 | ### Requirements 6 | * `tensorflow >= v1.3` 7 | * `numpy >= 1.13` 8 | * Optional requirements for IO include `h5py`, `larcv` 9 | 10 | ### Help 11 | An executable script can be found at `bin/dgcnn.py`. The script takes `train` or `inference` arguments. Try `--help` to list available arguments: 12 | ``` 13 | bin/dgcnn.py train --help 14 | ``` 15 | ### How to run 16 | Below is an example of how to train the network using `mydata.hdf5` data file with `hdf5` format, 4 GPUs with batch size 24 and mini-batch size of 6, store snapshot every 500 iterations, print out info (loss,accuracy,etc) every 10 iterations, and store tensorboard summary every 50 iterations. 17 | ``` 18 | bin/dgcnn.py train --gpus 0,1,2,3 -bs 24 -mbs 6 -chks 500 -rs 10 -ss 50 -if mydata.hdf5 -io h5 19 | ``` 20 | See `--help` to find more flags and a descipriton for arguments. 21 | 22 | 23 | 24 | -------------------------------------------------------------------------------- /bin/dgcnn.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | import os,sys 3 | DGCNN_DIR = os.path.dirname(os.path.abspath(__file__)) 4 | DGCNN_DIR = os.path.dirname(DGCNN_DIR) 5 | sys.path.insert(0,DGCNN_DIR) 6 | from dgcnn import DGCNN_FLAGS 7 | 8 | def main(): 9 | flags = DGCNN_FLAGS() 10 | flags.parse_args() 11 | 12 | if __name__ == '__main__': 13 | main() 14 | 15 | -------------------------------------------------------------------------------- /dgcnn/__init__.py: -------------------------------------------------------------------------------- 1 | from iotool import io_factory 2 | from model import build 3 | from trainval import trainval 4 | import ops 5 | from flags import DGCNN_FLAGS 6 | -------------------------------------------------------------------------------- /dgcnn/flags.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import argparse, os, sys 4 | from main_funcs import train, iotest, inference 5 | from distutils.util import strtobool 6 | class DGCNN_FLAGS: 7 | 8 | # flags for model 9 | NUM_CLASS = 2 10 | MODEL_NAME = "dgcnn" 11 | TRAIN = True 12 | KVALUE = 20 13 | DEBUG = True 14 | EDGE_CONV_LAYERS = 3 15 | EDGE_CONV_FILTERS = 64 16 | FC_LAYERS = 2 17 | FC_FILTERS = '512,256' 18 | 19 | # flags for train/inference 20 | SEED = -1 21 | LEARNING_RATE = 0.001 22 | GPUS = [0] 23 | MINIBATCH_SIZE = 1 24 | WEIGHT_PREFIX = './weights/snapshot' 25 | KVALUE = 20 26 | NUM_POINT = 2048 27 | NUM_CHANNEL = -1 28 | ITERATION = 10000 29 | REPORT_STEP = 100 30 | SUMMARY_STEP = 20 31 | CHECKPOINT_STEP = 500 32 | CHECKPOINT_NUM = 10 33 | CHECKPOINT_HOUR = 0.4 34 | 35 | # flags for IO 36 | IO_TYPE = 'h5' 37 | INPUT_FILE = '/scratch/kterao/dlprod_ppn_v08/dgcnn_p02_test_4ch.hdf5' 38 | OUTPUT_FILE = '' 39 | BATCH_SIZE = 1 40 | LOG_DIR = '' 41 | MODEL_PATH = '' 42 | DATA_KEY = 'data' 43 | LABEL_KEY = 'label' 44 | WEIGHT_KEY = '' 45 | SHUFFLE = 1 46 | 47 | def __init__(self): 48 | self._build_parsers() 49 | 50 | def _attach_common_args(self,parser): 51 | parser.add_argument('-kv','--kvalue',type=int,default=self.KVALUE,help='K value [default: %s]' % self.KVALUE) 52 | parser.add_argument('-db','--debug',type=strtobool,default=self.DEBUG, 53 | help='Extra verbose mode for debugging [default: %s]' % self.DEBUG) 54 | parser.add_argument('-ld','--log_dir', default=self.LOG_DIR, 55 | help='Log dir [default: %s]' % self.LOG_DIR) 56 | parser.add_argument('-sh','--shuffle',type=strtobool,default=self.SHUFFLE, 57 | help='Shuffle the data entries [default: %s]' % self.SHUFFLE) 58 | parser.add_argument('--gpus', type=str, default='0', 59 | help='GPUs to utilize (comma-separated integers') 60 | parser.add_argument('-ecl','--edge_conv_layers',type=int, default=self.EDGE_CONV_LAYERS, 61 | help='Number of edge-convolution layers [default: %s]' % self.EDGE_CONV_LAYERS) 62 | parser.add_argument('-ecf','--edge_conv_filters',type=str, default=str(self.EDGE_CONV_FILTERS), 63 | help='Number of filters in edge-convolution layers [default: %s]' % self.EDGE_CONV_FILTERS) 64 | parser.add_argument('-fcl','--fc_layers',type=int, default=self.FC_LAYERS, 65 | help='Number of fully-connected layers [default: %s]' % self.FC_LAYERS) 66 | parser.add_argument('-fcf','--fc_filters',type=str, default=str(self.FC_FILTERS), 67 | help='Number of filters in fully-connected layers [default: %s]' % self.FC_FILTERS) 68 | parser.add_argument('-nc','--num_class', type=int, default=self.NUM_CLASS, 69 | help='Number of classes [default: %s]' % self.NUM_CLASS) 70 | parser.add_argument('-np','--num_point', type=int, default=self.NUM_POINT, 71 | help='Point number [default: %s]' % self.NUM_POINT) 72 | parser.add_argument('-it','--iteration', type=int, default=self.ITERATION, 73 | help='Iteration to run [default: %s]' % self.ITERATION) 74 | parser.add_argument('-bs','--batch_size', type=int, default=self.BATCH_SIZE, 75 | help='Batch Size during training for updating weights [default: %s]' % self.BATCH_SIZE) 76 | parser.add_argument('-mbs','--minibatch_size', type=int, default=self.MINIBATCH_SIZE, 77 | help='Mini-Batch Size during training for each GPU [default: %s]' % self.MINIBATCH_SIZE) 78 | parser.add_argument('-rs','--report_step', type=int, default=self.REPORT_STEP, 79 | help='Period (in steps) to print out loss and accuracy [default: %s]' % self.REPORT_STEP) 80 | parser.add_argument('-mn','--model_name', type=str, default=self.MODEL_NAME, 81 | help='model name identifier [default: %s]' % self.MODEL_NAME) 82 | parser.add_argument('-mp','--model_path', type=str, default=self.MODEL_PATH, 83 | help='model checkpoint file path [default: %s]' % self.MODEL_PATH) 84 | parser.add_argument('-io','--io_type',type=str,default=self.IO_TYPE, 85 | help='IO handler type [default: %s]' % self.IO_TYPE) 86 | parser.add_argument('-if','--input_file',type=str,default=self.INPUT_FILE, 87 | help='comma-separated input file list [default: %s]' % self.INPUT_FILE) 88 | parser.add_argument('-of','--output_file',type=str,default=self.OUTPUT_FILE, 89 | help='output file name [default: %s]' % self.OUTPUT_FILE) 90 | parser.add_argument('-dkey','--data_key',type=str,default=self.DATA_KEY, 91 | help='A keyword to fetch data from file [default: %s]' % self.DATA_KEY) 92 | parser.add_argument('-lkey','--label_key',type=str,default=self.LABEL_KEY, 93 | help='A keyword to fetch label from file [default: %s]' % self.LABEL_KEY) 94 | return parser 95 | 96 | def _build_parsers(self): 97 | 98 | self.parser = argparse.ArgumentParser(description="Edge-GCNN Configuration Flags") 99 | subparsers = self.parser.add_subparsers(title="Modules", description="Valid subcommands", dest='script', help="aho") 100 | 101 | # train parser 102 | train_parser = subparsers.add_parser("train", help="Train Edge-GCNN") 103 | train_parser.add_argument('-sd','--seed', default=self.SEED, 104 | help='Seed for random number generators [default: %s]' % self.SEED) 105 | train_parser.add_argument('-wp','--weight_prefix', default=self.WEIGHT_PREFIX, 106 | help='Prefix (directory + file prefix) for snapshots of weights [default: %s]' % self.WEIGHT_PREFIX) 107 | train_parser.add_argument('-lr','--learning_rate', type=float, default=self.LEARNING_RATE, 108 | help='Initial learning rate [default: %s]' % self.LEARNING_RATE) 109 | train_parser.add_argument('-ss','--summary_step', type=int, default=self.SUMMARY_STEP, 110 | help='Period (in steps) to store summary in tensorboard log [default: %s]' % self.SUMMARY_STEP) 111 | train_parser.add_argument('-chks','--checkpoint_step', type=int, default=self.CHECKPOINT_STEP, 112 | help='Period (in steps) to store snapshot of weights [default: %s]' % self.CHECKPOINT_STEP) 113 | train_parser.add_argument('-chkn','--checkpoint_num', type=int, default=self.CHECKPOINT_NUM, 114 | help='Number of the latest checkpoint to keep [default: %s]' % self.CHECKPOINT_NUM) 115 | train_parser.add_argument('-chkh','--checkpoint_hour', type=float, default=self.CHECKPOINT_HOUR, 116 | help='Period (in hours) to store checkpoint [default: %s]' % self.CHECKPOINT_HOUR) 117 | train_parser.add_argument('-wkey','--weight_key',type=str,default=self.WEIGHT_KEY, 118 | help='A keyword to fetch weight from file [default: %s]' % self.WEIGHT_KEY) 119 | 120 | # inference parser 121 | inference_parser = subparsers.add_parser("inference",help="Run inference of Edge-GCNN") 122 | # IO test parser 123 | iotest_parser = subparsers.add_parser("iotest", help="Test iotools for Edge-GCNN") 124 | iotest_parser.add_argument('-wkey','--weight_key',type=str,default=self.WEIGHT_KEY, 125 | help='A keyword to fetch weight from file [default: %s]' % self.WEIGHT_KEY) 126 | 127 | # attach common parsers 128 | self.train_parser = self._attach_common_args(train_parser) 129 | self.inference_parser = self._attach_common_args(inference_parser) 130 | self.iotest_parser = self._attach_common_args(iotest_parser) 131 | 132 | # attach executables 133 | self.train_parser.set_defaults(func=train) 134 | self.inference_parser.set_defaults(func=inference) 135 | self.iotest_parser.set_defaults(func=iotest) 136 | def parse_args(self): 137 | args = self.parser.parse_args() 138 | self.update(vars(args)) 139 | print("\n\n-- CONFIG --") 140 | for name in vars(self): 141 | attribute = getattr(self,name) 142 | if type(attribute) == type(self.parser): continue 143 | print("%s = %r" % (name, getattr(self, name))) 144 | 145 | # Set random seed for reproducibility 146 | np.random.seed(self.SEED) 147 | tf.set_random_seed(self.SEED) 148 | args.func(self) 149 | 150 | def update(self, args): 151 | for name,value in args.iteritems(): 152 | if name in ['func','script']: continue 153 | setattr(self, name.upper(), args[name]) 154 | os.environ['CUDA_VISIBLE_DEVICES']=self.GPUS 155 | self.GPUS=[int(gpu) for gpu in self.GPUS.split(',')] 156 | self.INPUT_FILE=[str(f) for f in self.INPUT_FILE.split(',')] 157 | if self.EDGE_CONV_FILTERS.find(',')>0: 158 | self.EDGE_CONV_FILTERS = [int(v) for v in self.EDGE_CONV_FILTERS.split(',')] 159 | else: 160 | self.EDGE_CONV_FILTERS = int(self.EDGE_CONV_FILTERS) 161 | if self.FC_FILTERS.find(',')>0: 162 | self.FC_FILTERS = [int(v) for v in self.FC_FILTERS.split(',')] 163 | else: 164 | self.FC_FILTERS = int(self.FC_FILTERS) 165 | if self.SEED < 0: 166 | import time 167 | self.SEED = int(time.time()) 168 | 169 | if __name__ == '__main__': 170 | flags=GCNN_FLAGS() 171 | flags.parse_args() 172 | 173 | -------------------------------------------------------------------------------- /dgcnn/iotool.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import numpy as np 5 | import sys 6 | class io_base(object): 7 | 8 | def __init__(self,flags): 9 | self._batch_size = flags.BATCH_SIZE 10 | self._num_entries = -1 11 | self._num_channels = -1 12 | 13 | def batch_size(self,size=None): 14 | if size is None: return self._batch_size 15 | self._batch_size = int(size) 16 | 17 | def num_entries(self): 18 | return self._num_entries 19 | 20 | def num_channels(self): 21 | return self._num_channels 22 | 23 | def initialize(self): 24 | raise NotImplementedError 25 | 26 | def store(self,idx,softmax): 27 | raise NotImplementedError 28 | 29 | def next(self): 30 | raise NotImplementedError 31 | 32 | def finalize(self): 33 | raise NotImplementedError 34 | 35 | class io_larcv(io_base): 36 | 37 | def __init__(self,flags): 38 | super(io_larcv,self).__init__(flags=flags) 39 | self._flags = flags 40 | self._data = None 41 | self._label = None 42 | self._weight = None 43 | self._fout = None 44 | self._last_entry = -1 45 | self._event_keys = [] 46 | self._metas = [] 47 | 48 | def initialize(self): 49 | self._last_entry = -1 50 | self._event_keys = [] 51 | self._metas = [] 52 | # configure the input 53 | from larcv import larcv 54 | from ROOT import TChain 55 | ch_data = TChain('sparse3d_%s_tree' % self._flags.DATA_KEY) 56 | ch_label = None 57 | ch_weight = None 58 | if self._flags.LABEL_KEY: 59 | ch_label = TChain('sparse3d_%s_tree' % self._flags.LABEL_KEY) 60 | if self._flags.WEIGHT_KEY: 61 | ch_weight = TChain('sparse3d_%s_tree' % self._flags.WEIGHT_KEY) 62 | for f in self._flags.INPUT_FILE: 63 | ch_data.AddFile(f) 64 | if ch_label: ch_label.AddFile(f) 65 | if ch_weight: ch_weight.AddFile(f) 66 | self._data = [] 67 | self._label = [] 68 | self._weight = [] 69 | br_data,br_label,br_weight=(None,None,None) 70 | event_fraction = 1./ch_data.GetEntries() * 100. 71 | total_point = 0. 72 | for i in range(ch_data.GetEntries()): 73 | ch_data.GetEntry(i) 74 | if ch_label: ch_label.GetEntry(i) 75 | if ch_weight: ch_weight.GetEntry(i) 76 | if br_data is None: 77 | br_data = getattr(ch_data, 'sparse3d_%s_branch' % self._flags.DATA_KEY) 78 | if ch_label: br_label = getattr(ch_label, 'sparse3d_%s_branch' % self._flags.LABEL_KEY) 79 | if ch_weight: br_weight = getattr(ch_weight,'sparse3d_%s_branch' % self._flags.WEIGHT_KEY) 80 | num_point = br_data.as_vector().size() 81 | if num_point < 256: continue 82 | np_data = np.zeros(shape=(num_point,4),dtype=np.float32) 83 | larcv.fill_3d_pcloud(br_data, np_data) 84 | self._data.append(np_data) 85 | self._event_keys.append((br_data.run(),br_data.subrun(),br_data.event())) 86 | self._metas.append(larcv.Voxel3DMeta(br_data.meta())) 87 | if ch_label: 88 | np_label = np.zeros(shape=(num_point,1),dtype=np.float32) 89 | larcv.fill_3d_pcloud(br_label, np_label) 90 | np_label = np_label.reshape([num_point]) - 1. 91 | self._label.append(np_label) 92 | if ch_weight: 93 | np_weight = np.zeros(shape=(num_point,1),dtype=np.float32) 94 | larcv.fill_3d_pcloud(br_weight, np_weight) 95 | np_weight = np_weight.reshape([num_point]) 96 | np_weight = np_weight / np_weight.sum() * len(np_weight) 97 | self._weight.append(np_weight) 98 | total_point += np_data.size 99 | sys.stdout.write('Processed %d%% ... %d MB\r' % (int(event_fraction*i),int(total_point*4*2/1.e6))) 100 | sys.stdout.flush() 101 | 102 | sys.stdout.write('\n') 103 | sys.stdout.flush() 104 | self._num_channels = self._data[-1].shape[-1] 105 | self._num_entries = len(self._data) 106 | # Output 107 | if self._flags.OUTPUT_FILE: 108 | import tempfile 109 | cfg = ''' 110 | IOManager: { 111 | Verbosity: 2 112 | Name: "IOManager" 113 | IOMode: 1 114 | OutFileName: "%s" 115 | InputFiles: [] 116 | InputDirs: [] 117 | StoreOnlyType: [] 118 | StoreOnlyName: [] 119 | } 120 | ''' 121 | cfg = cfg % self._flags.OUTPUT_FILE 122 | cfg_file = tempfile.NamedTemporaryFile('w') 123 | cfg_file.write(cfg) 124 | cfg_file.flush() 125 | self._fout = larcv.IOManager(cfg_file.name) 126 | self._fout.initialize() 127 | 128 | def next(self): 129 | data,label,weight=(None,None,None) 130 | start,end=(-1,-1) 131 | if self._flags.SHUFFLE: 132 | start = int(np.random.random() * (self.num_entries() - self.batch_size())) 133 | end = start + self.batch_size() 134 | idx = np.arange(start,end,1) 135 | data = self._data[start:end] 136 | if len(self._label) > 0: label = self._label[start:end] 137 | if len(self._weight) > 0: weight = self._weight[start:end] 138 | else: 139 | start = self._last_entry+1 140 | end = start + self.batch_size() 141 | if end < self.num_entries(): 142 | idx = np.arange(start,end,1) 143 | data = self._data[start:end] 144 | if len(self._label) > 0: label = self._label[start:end] 145 | if len(self._weight) > 0: weight = self._weight[start:end] 146 | else: 147 | idx = np.concatenate([np.arange(start,self.num_entries(),1),np.arange(0,end-self.num_entries(),1)]) 148 | data = self._data[start:] + self._data[0:end-self.num_entries()] 149 | if len(self._label) > 0: label = self._label[start:] + self._label[0:end-self.num_entries()] 150 | if len(self._weight) > 0: weight = self._weight[start:] + self._weight[0:end-self._num_entries()] 151 | self._last_entry = idx[-1] 152 | 153 | return idx,data,label,weight 154 | 155 | def store(self,idx,softmax): 156 | from larcv import larcv 157 | if self._fout is None: 158 | raise NotImplementedError 159 | idx=int(idx) 160 | if idx >= self.num_entries(): 161 | raise ValueError 162 | keys = self._event_keys[idx] 163 | meta = self._metas[idx] 164 | 165 | larcv_data = self._fout.get_data('sparse3d',self._flags.DATA_KEY) 166 | data = self._data[idx] 167 | vs = larcv.as_tensor3d(data,meta,0.) 168 | larcv_data.set(vs,meta) 169 | 170 | pos = data[:,0:3] 171 | score = np.max(softmax,axis=1).reshape([len(softmax),1]) 172 | score = np.concatenate([pos,score],axis=1) 173 | prediction = np.argmax(softmax,axis=1).astype(np.float32).reshape([len(softmax),1]) 174 | prediction = np.concatenate([pos,prediction],axis=1) 175 | 176 | larcv_softmax = self._fout.get_data('sparse3d','softmax') 177 | vs = larcv.as_tensor3d(score,meta,-1.) 178 | larcv_softmax.set(vs,meta) 179 | 180 | larcv_prediction = self._fout.get_data('sparse3d','prediction') 181 | vs = larcv.as_tensor3d(prediction,meta,-1.) 182 | larcv_prediction.set(vs,meta) 183 | 184 | if len(self._label) > 0: 185 | label = self._label[idx] 186 | label = label.astype(np.float32).reshape([len(label),1]) 187 | label = np.concatenate([pos,label],axis=1) 188 | larcv_label = self._fout.get_data('sparse3d','label') 189 | vs = larcv.as_tensor3d(label,meta,-1.) 190 | larcv_label.set(vs,meta) 191 | 192 | self._fout.set_id(keys[0],keys[1],keys[2]) 193 | self._fout.save_entry() 194 | 195 | def finalize(self): 196 | if self._fout: 197 | self._fout.finalize() 198 | 199 | class io_h5(io_base): 200 | 201 | def __init__(self,flags): 202 | super(io_h5,self).__init__(flags=flags) 203 | self._flags = flags 204 | self._data = None 205 | self._label = None 206 | self._weight = None 207 | self._fout = None 208 | self._ohandler_data = None 209 | self._ohandler_label = None 210 | self._ohandler_softmax = None 211 | self._has_label = False 212 | 213 | def initialize(self): 214 | self._last_entry = -1 215 | # Prepare input 216 | import h5py as h5 217 | self._data = None 218 | self._label = None 219 | self._weight = None 220 | for f in self._flags.INPUT_FILE: 221 | f = h5.File(f,'r') 222 | if self._data is None: 223 | self._data = np.array(f[self._flags.DATA_KEY ]) 224 | if self._flags.LABEL_KEY : self._label = np.array(f[self._flags.LABEL_KEY]) 225 | if self._flags.WEIGHT_KEY: self._weight = np.array(f[self._flags.WEIGHT_KEY]) 226 | else: 227 | self._data = np.concatenate(self._data, np.array(f[self._flags.DATA_KEY ])) 228 | if self._label : self._label = np.concatenate(self._label, np.array(f[self._flags.LABEL_KEY ])) 229 | if self._weight : self._weight = np.concatenate(self._weight,np.array(f[self._flags.WEIGHT_KEY])) 230 | self._num_channels = self._data[-1].shape[-1] 231 | self._num_entries = len(self._data) 232 | # Prepare output 233 | if self._flags.OUTPUT_FILE: 234 | import tables 235 | FILTERS = tables.Filters(complib='zlib', complevel=5) 236 | self._fout = tables.open_file(self._flags.OUTPUT_FILE,mode='w', filters=FILTERS) 237 | data_shape = list(self._data[0].shape) 238 | data_shape.insert(0,0) 239 | self._ohandler_data = self._fout.create_earray(self._fout.root,self._flags.DATA_KEY,tables.Float32Atom(),shape=data_shape) 240 | self._ohandler_softmax = self._fout.create_earray(self._fout.root,'softmax',tables.Float32Atom(),shape=data_shape) 241 | if self._label: 242 | data_shape = list(self._label[0].shape) 243 | data_shape.insert(0,0) 244 | self._ohandler_label = self._fout.create_earray(self._fout.root,self._flags.LABEL_KEY,tables.Float32Atom(),shape=data_shape) 245 | def store(self,idx,softmax): 246 | if self._fout is None: 247 | raise NotImplementedError 248 | idx=int(idx) 249 | if idx >= self.num_entries(): 250 | raise ValueError 251 | data = self._data[idx] 252 | self._ohandler_data.append(data[None]) 253 | self._ohandler_softmax.append(softmax[None]) 254 | if self._label is not None: 255 | label = self._label[idx] 256 | self._ohandler_label.append(label[None]) 257 | 258 | def next(self): 259 | idx = None 260 | if self._flags.SHUFFLE: 261 | idx = np.arange(self.num_entries()) 262 | np.random.shuffle(idx) 263 | idx = idx[0:self.batch_size()] 264 | else: 265 | start = self._last_entry+1 266 | end = start + self.batch_size() 267 | if end < self.num_entries(): 268 | idx = np.arange(start,end) 269 | else: 270 | idx = np.concatenate([np.arange(start,self.num_entries()),np.arange(0,end-self.num_entries())]) 271 | self._last_entry = idx[-1] 272 | data = self._data[idx, ...] 273 | label,weight=(None,None) 274 | if self._label : label = self._label[idx, ...] 275 | if self._weight : weight = self._weight[idx, ...] 276 | return idx, data, label, weight 277 | 278 | def finalize(self): 279 | if self._fout: 280 | self._fout.close() 281 | 282 | def io_factory(flags): 283 | if flags.IO_TYPE == 'h5': 284 | return io_h5(flags) 285 | if flags.IO_TYPE == 'larcv': 286 | return io_larcv(flags) 287 | raise NotImplementedError 288 | -------------------------------------------------------------------------------- /dgcnn/main_funcs.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import os,time,datetime,sys 5 | import numpy as np 6 | import dgcnn 7 | import tensorflow as tf 8 | 9 | def round_decimals(val,digits): 10 | factor = float(np.power(10,digits)) 11 | return int(val * factor+0.5) / factor 12 | 13 | def iteration_from_filename(file_name): 14 | return int((file_name.split('-'))[-1]) 15 | 16 | def iotest(flags): 17 | # IO configuration 18 | io = dgcnn.io_factory(flags) 19 | io.initialize() 20 | num_entries = io.num_entries() 21 | ctr = 0 22 | while ctr < num_entries: 23 | idx,data,label,weight=io.next() 24 | msg = str(ctr) + '/' + str(num_entries) + ' ... ' + str(idx) + ' ' + str(data[0].shape) 25 | if label: 26 | msg += str(label[0].shape) 27 | if weight: 28 | msg += str(weight[0].shape) 29 | print(msg) 30 | ctr += len(data) 31 | io.finalize() 32 | 33 | class Handlers: 34 | sess = None 35 | data_io = None 36 | csv_logger = None 37 | weight_io = None 38 | train_logger = None 39 | iteration = 0 40 | 41 | def train(flags): 42 | 43 | flags.TRAIN = True 44 | handlers = prepare(flags) 45 | train_loop(flags,handlers) 46 | 47 | def inference(flags): 48 | 49 | flags.TRAIN = False 50 | handlers = prepare(flags) 51 | inference_loop(flags,handlers) 52 | 53 | def prepare(flags): 54 | 55 | handlers = Handlers() 56 | # assert 57 | if flags.BATCH_SIZE % (flags.MINIBATCH_SIZE * len(flags.GPUS)): 58 | msg = '--batch_size (%d) must be a modular of --gpus (%d) * --minibatch_size (%d)\n' 59 | msg = msg % (flags.BATCH_SIZE,flags.MINIBATCH_SIZE,len(flags.GPUS)) 60 | sys.stderr.write(msg) 61 | sys.exit(1) 62 | 63 | # IO configuration 64 | handlers.data_io = dgcnn.io_factory(flags) 65 | handlers.data_io.initialize() 66 | _,train_data,_,_ = handlers.data_io.next() 67 | 68 | # Trainer configuration 69 | flags.NUM_CHANNEL = handlers.data_io.num_channels() 70 | handlers.trainer = dgcnn.trainval(flags) 71 | handlers.trainer.initialize() 72 | 73 | # Create a session 74 | config = tf.ConfigProto() 75 | config.gpu_options.allow_growth = True 76 | config.allow_soft_placement = True 77 | handlers.sess = tf.Session(config=config) 78 | init = tf.group(tf.global_variables_initializer(), 79 | tf.local_variables_initializer()) 80 | handlers.sess.run(init) 81 | 82 | handlers.weight_io = tf.train.Saver(max_to_keep=flags.CHECKPOINT_NUM, 83 | keep_checkpoint_every_n_hours=flags.CHECKPOINT_HOUR) 84 | if flags.WEIGHT_PREFIX: 85 | save_dir = flags.WEIGHT_PREFIX[0:flags.WEIGHT_PREFIX.rfind('/')] 86 | if save_dir and not os.path.isdir(save_dir): os.makedirs(save_dir) 87 | 88 | handlers.iteration = 0 89 | loaded_iteration = 0 90 | if flags.MODEL_PATH: 91 | handlers.weight_io.restore(handlers.sess, flags.MODEL_PATH) 92 | loaded_iteration = iteration_from_filename(flags.MODEL_PATH) 93 | if flags.TRAIN: handlers.iteration = loaded_iteration+1 94 | 95 | if flags.LOG_DIR: 96 | if not os.path.exists(flags.LOG_DIR): os.mkdir(flags.LOG_DIR) 97 | handlers.train_logger = tf.summary.FileWriter(flags.LOG_DIR) 98 | handlers.train_logger.add_graph(handlers.sess.graph) 99 | logname = '%s/train_log-%07d.csv' % (flags.LOG_DIR,loaded_iteration) 100 | if not flags.TRAIN: 101 | logname = '%s/inference_log-%07d.csv' % (flags.LOG_DIR,loaded_iteration) 102 | handlers.csv_logger = open(logname,'w') 103 | 104 | return handlers 105 | 106 | def train_loop(flags,handlers): 107 | 108 | handlers.csv_logger.write('iter,epoch') 109 | handlers.csv_logger.write(',titer,ttrain,tio,tsave,tsummary') 110 | handlers.csv_logger.write(',tsumiter,tsumtrain,tsumio,tsumsave,tsumsummary') 111 | handlers.csv_logger.write(',loss,accuracy\n') 112 | 113 | tsum = 0. 114 | tsum_train = 0. 115 | tsum_io = 0. 116 | tsum_save = 0. 117 | tsum_summary = 0. 118 | while handlers.iteration < flags.ITERATION: 119 | 120 | tstamp_iteration = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S') 121 | tstart_iteration = time.time() 122 | 123 | report_step = flags.REPORT_STEP and ((handlers.iteration+1) % flags.REPORT_STEP == 0) 124 | summary_step = flags.SUMMARY_STEP and handlers.train_logger and ((handlers.iteration+1) % flags.SUMMARY_STEP == 0) 125 | checkpt_step = flags.CHECKPOINT_STEP and flags.WEIGHT_PREFIX and ((handlers.iteration+1) % flags.CHECKPOINT_STEP == 0) 126 | 127 | tstart = time.time() 128 | idx,data,label,weight = handlers.data_io.next() 129 | tspent_io = time.time() - tstart 130 | tsum_io += tspent_io 131 | 132 | current_idx = 0 133 | loss_v = [] 134 | accuracy_v = [] 135 | handlers.trainer.zero_gradients(handlers.sess) 136 | # Accummulate gradients 137 | tspent_train = 0. 138 | tspent_summary = 0. 139 | while current_idx < flags.BATCH_SIZE: 140 | tstart = time.time() 141 | data_v = [] 142 | label_v = [] 143 | weight_v = None 144 | if weight is not None: weight_v = [] 145 | for _ in flags.GPUS: 146 | start = current_idx 147 | end = current_idx + flags.MINIBATCH_SIZE 148 | data_v.append(data[start:end]) 149 | label_v.append(label[start:end]) 150 | if weight is not None: 151 | weight_v.append(weight[start:end]) 152 | current_idx = end 153 | # compute gradients 154 | make_summary = summary_step and (current_idx == flags.BATCH_SIZE) 155 | res = handlers.trainer.accum_gradient(handlers.sess,data_v,label_v,weight_v,summary=make_summary) 156 | accuracy_v.append(res[1]) 157 | loss_v.append(res[2]) 158 | tspent_train = tspent_train + (time.time() - tstart) 159 | # log summary 160 | if make_summary: 161 | tstart = time.time() 162 | handlers.train_logger.add_summary(res[3],handlers.iteration) 163 | tspent_summary = time.time() - tstart 164 | # Apply gradients 165 | tstart = time.time() 166 | handlers.trainer.apply_gradient(handlers.sess) 167 | tspent_train = tspent_train + (time.time() - tstart) 168 | 169 | tsum_train += tspent_train 170 | tsum_summary += tspent_summary 171 | 172 | # Compute loss/accuracy 173 | loss = np.mean(loss_v) 174 | accuracy = np.mean(accuracy_v) 175 | epoch = handlers.iteration * float(flags.BATCH_SIZE) / handlers.data_io.num_entries() 176 | # Save snapshot 177 | tspent_save = 0. 178 | if checkpt_step: 179 | tstart = time.time() 180 | ssf_path = handlers.weight_io.save(handlers.sess,flags.WEIGHT_PREFIX,global_step=handlers.iteration) 181 | tspent_save = time.time() - tstart 182 | print('saved @',ssf_path) 183 | # Report (logger) 184 | if handlers.csv_logger: 185 | tspent_iteration = time.time() - tstart_iteration 186 | tsum += tspent_iteration 187 | csv_data = '%d,%g,' % (handlers.iteration,epoch) 188 | csv_data += '%g,%g,%g,%g,%g,' % (tspent_iteration,tspent_train,tspent_io,tspent_save,tspent_summary) 189 | csv_data += '%g,%g,%g,%g,%g,' % (tsum,tsum_train,tsum_io,tsum_save,tsum_summary) 190 | csv_data += '%g,%g\n' % (loss,accuracy) 191 | handlers.csv_logger.write(csv_data) 192 | # Report (stdout) 193 | if report_step: 194 | loss = round_decimals(loss,4) 195 | accuracy = round_decimals(accuracy,4) 196 | tfrac = round_decimals(tspent_train/tspent_iteration*100.,2) 197 | epoch = round_decimals(epoch,2) 198 | mem = handlers.sess.run(tf.contrib.memory_stats.MaxBytesInUse()) 199 | msg = 'Iteration %d (epoch %g) @ %s ... train time fraction %g%% max mem. %g ... loss %g accuracy %g' 200 | msg = msg % (handlers.iteration,epoch,tstamp_iteration,tfrac,mem,loss,accuracy) 201 | print(msg) 202 | sys.stdout.flush() 203 | if handlers.csv_logger: handlers.csv_logger.flush() 204 | if handlers.train_logger: handlers.train_logger.flush() 205 | # Increment iteration counter 206 | handlers.iteration +=1 207 | 208 | handlers.train_logger.close() 209 | handlers.csv_logger.close() 210 | handlers.data_io.finalize() 211 | 212 | def inference_loop(flags,handlers): 213 | handlers.csv_logger.write('iter,epoch') 214 | handlers.csv_logger.write(',titer,tinference,tio') 215 | handlers.csv_logger.write(',tsumiter,tsuminference,tsumio') 216 | handlers.csv_logger.write(',loss,accuracy\n') 217 | tsum = 0. 218 | tsum_io = 0. 219 | tsum_inference = 0. 220 | while handlers.iteration < flags.ITERATION: 221 | 222 | tstamp_iteration = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S') 223 | tstart_iteration = time.time() 224 | 225 | report_step = flags.REPORT_STEP and ((handlers.iteration+1) % flags.REPORT_STEP == 0) 226 | 227 | tstart = time.time() 228 | idx,data,label,weight = handlers.data_io.next() 229 | tspent_io = time.time() - tstart 230 | tsum_io += tspent_io 231 | 232 | current_idx = 0 233 | softmax_vv = [] 234 | loss_v = [] 235 | accuracy_v = [] 236 | 237 | # Run inference 238 | tspent_inference = 0. 239 | tstart = time.time() 240 | while current_idx < flags.BATCH_SIZE: 241 | data_v = [] 242 | label_v = None 243 | weight_v = None 244 | if label is not None: label_v = [] 245 | if weight is not None: weight_v = [] 246 | for _ in flags.GPUS: 247 | start = current_idx 248 | end = current_idx + flags.MINIBATCH_SIZE 249 | data_v.append(data[start:end]) 250 | if label is not None: 251 | label_v.append(label[start:end]) 252 | if weight is not None: 253 | weight_v.append(weight[start:end]) 254 | current_idx = end 255 | # compute gradients 256 | res = handlers.trainer.inference(handlers.sess,data_v,label_v,weight_v) 257 | if flags.LABEL_KEY: 258 | softmax_vv = softmax_vv + res[0:-2] 259 | accuracy_v.append(res[-2]) 260 | loss_v.append(res[-1]) 261 | else: 262 | softmax_vv = softmax_vv + res 263 | tspent_inference = tspent_inference + (time.time() - tstart) 264 | tsum_inference += tspent_inference 265 | 266 | # Store output if requested 267 | if flags.OUTPUT_FILE: 268 | idx_ctr = 0 269 | for softmax_v in softmax_vv: 270 | for softmax in softmax_v: 271 | handlers.data_io.store(idx[idx_ctr],softmax) 272 | idx_ctr += 1 273 | 274 | # Compute loss/accuracy 275 | loss,accuracy=[-1,-1] 276 | if flags.LABEL_KEY: 277 | loss = np.mean(loss_v) 278 | accuracy = np.mean(accuracy_v) 279 | epoch = handlers.iteration * float(flags.BATCH_SIZE) / handlers.data_io.num_entries() 280 | # Report (logger) 281 | if handlers.csv_logger: 282 | tspent_iteration = time.time() - tstart_iteration 283 | tsum += tspent_iteration 284 | csv_data = '%d,%g,' % (handlers.iteration,epoch) 285 | csv_data += '%g,%g,%g,' % (tspent_iteration,tspent_inference,tspent_io) 286 | csv_data += '%g,%g,%g,' % (tsum,tsum_inference,tsum_io) 287 | csv_data += '%g,%g\n' % (loss,accuracy) 288 | handlers.csv_logger.write(csv_data) 289 | # Report (stdout) 290 | if report_step: 291 | loss = round_decimals(loss,4) 292 | accuracy = round_decimals(accuracy,4) 293 | tfrac = round_decimals(tspent_inference/tspent_iteration*100.,2) 294 | epoch = round_decimals(epoch,2) 295 | mem = handlers.sess.run(tf.contrib.memory_stats.MaxBytesInUse()) 296 | msg = 'Iteration %d (epoch %g) @ %s ... inference time fraction %g%% max mem. %g ... loss %g accuracy %g' 297 | msg = msg % (handlers.iteration,epoch,tstamp_iteration,tfrac,mem,loss,accuracy) 298 | print(msg) 299 | sys.stdout.flush() 300 | if handlers.csv_logger: handlers.csv_logger.flush() 301 | # Increment iteration counter 302 | handlers.iteration +=1 303 | 304 | handlers.csv_logger.close() 305 | handlers.data_io.finalize() 306 | 307 | -------------------------------------------------------------------------------- /dgcnn/model.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import tensorflow.python.platform 5 | import tensorflow as tf 6 | import tensorflow.contrib.slim as slim 7 | import dgcnn 8 | 9 | def build(point_cloud, flags): 10 | 11 | num_edge_conv = int(flags.EDGE_CONV_LAYERS) 12 | num_edge_filters = flags.EDGE_CONV_FILTERS 13 | num_fc = int(flags.FC_LAYERS) 14 | num_fc_filters = flags.FC_FILTERS 15 | is_training = bool(flags.TRAIN) 16 | k = int(flags.KVALUE) 17 | debug = bool(flags.DEBUG) 18 | num_class = int(flags.NUM_CLASS) 19 | 20 | net = point_cloud 21 | batch_size = tf.shape(net)[0] 22 | num_point = tf.shape(net)[1] 23 | if debug: 24 | print('\n') 25 | print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 26 | 27 | if flags.MODEL_NAME == 'dgcnn': 28 | tensors = dgcnn.ops.repeat_edge_conv(net, 29 | repeat=num_edge_conv, 30 | k=k, 31 | num_filters=num_edge_filters, 32 | trainable=is_training, 33 | debug=debug) 34 | elif flags.MODEL_NAME in ['residual-dgcnn','residual-dgcnn-nofc']: 35 | tensors = dgcnn.ops.repeat_residual_edge_conv(net, 36 | repeat=num_edge_conv, 37 | k=k, 38 | num_filters=num_edge_filters, 39 | trainable=is_training, 40 | debug=debug) 41 | else: 42 | print('Unsupported MODEL_NAME: %s' % flags.MODEL_NAME) 43 | raise NotImplementedError 44 | 45 | if flags.MODEL_NAME == 'residual-dgcnn-nofc': 46 | net = slim.conv2d(inputs = tensors[-1], 47 | num_outputs = num_class, 48 | kernel_size = 1, 49 | stride = 1, 50 | trainable = True, 51 | padding = 'VALID', 52 | normalizer_fn = slim.batch_norm, 53 | scope = 'Final') 54 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 55 | 56 | net = tf.squeeze(net, axis=-2) 57 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 58 | return net 59 | 60 | concat = [] 61 | for i in range(num_edge_conv): 62 | concat.append(tensors[3*i+2]) 63 | concat = tf.concat(concat,axis=-1) 64 | 65 | net = slim.conv2d(inputs = concat, 66 | num_outputs = 1024, 67 | kernel_size = 1, 68 | stride = 1, 69 | trainable = True, 70 | padding = 'VALID', 71 | normalizer_fn = slim.batch_norm, 72 | scope = 'MergedEdgeConv') 73 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 74 | tensors.append(net) 75 | 76 | from tensorflow.python.ops import gen_nn_ops 77 | net = gen_nn_ops.max_pool_v2(net, ksize=[1,num_point,1,1], strides=[1,1,1,1], padding='VALID', name='maxpool0') 78 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 79 | 80 | net = tf.reshape(net,[batch_size,-1,1,1024]) 81 | net = tf.tile(net, [1, num_point, 1, 1]) 82 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 83 | concat = [net] + tensors 84 | 85 | net = tf.concat(values=concat, axis=3) 86 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 87 | 88 | net = dgcnn.ops.fc(net=net, repeat=num_fc, num_filters=num_fc_filters, trainable=is_training, debug=debug) 89 | 90 | if is_training: 91 | net = tf.nn.dropout(net, 0.7, None) 92 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 93 | 94 | net = slim.conv2d(inputs = net, 95 | num_outputs = num_class, 96 | kernel_size = 1, 97 | stride = 1, 98 | trainable = True, 99 | padding = 'VALID', 100 | normalizer_fn = slim.batch_norm, 101 | scope = 'Final') 102 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 103 | 104 | net = tf.squeeze(net, axis=-2) 105 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 106 | return net 107 | 108 | -------------------------------------------------------------------------------- /dgcnn/ops.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import numpy as np 5 | import tensorflow as tf 6 | import tensorflow.contrib.layers as L 7 | import tensorflow.contrib.slim as slim 8 | def k_nn(points, k): 9 | # The shape of points (B, N, C) 10 | # Use linear algebra to find all pairwise distance in parallel 11 | M = points 12 | transpose = tf.transpose(M, perm=[0, 2, 1]) 13 | inner_prod = tf.matmul(M, transpose) 14 | squared = tf.reduce_sum(tf.square(M), axis=-1, keepdims=True) 15 | squared_tranpose = tf.transpose(squared, perm=[0, 2, 1]) 16 | nn_dist = squared + squared_tranpose - 2 * inner_prod 17 | # Next pick the top k shortest ones 18 | _, idx = tf.nn.top_k(-nn_dist,k=k) 19 | return idx 20 | 21 | def edges(points, k=20): 22 | # points (B, N, C) 23 | knn_idx = k_nn(points, k) 24 | 25 | points_central = points 26 | batch_size = tf.shape(points)[0] 27 | num_points = tf.shape(points)[1] 28 | num_dims = tf.shape(points)[2] 29 | 30 | idx_ = tf.range (batch_size) * num_points 31 | idx_ = tf.reshape (idx_, [batch_size, 1, 1]) 32 | 33 | points_flat = tf.reshape (points, [-1, num_dims]) 34 | points_neighbors = tf.gather (points_flat, knn_idx+idx_) 35 | points_central = tf.expand_dims (points_central, axis=-2) 36 | 37 | points_central = tf.tile (points_central, [1, 1, k, 1]) 38 | 39 | edge_feature = tf.concat([points_central, points_neighbors-points_central], axis=-1) 40 | return edge_feature 41 | 42 | def edge_conv(point_cloud, k, num_filters, trainable, activation=tf.nn.relu, debug=False): 43 | 44 | net = point_cloud 45 | net = edges(net, k=k) 46 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 47 | net = slim.conv2d(inputs = net, 48 | num_outputs = num_filters, 49 | kernel_size = 1, 50 | stride = 1, 51 | trainable = trainable, 52 | padding = 'VALID', 53 | normalizer_fn = slim.batch_norm, 54 | scope = 'conv0') 55 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 56 | net_max = tf.reduce_max (net, axis=-2, keepdims=True) 57 | net_mean = tf.reduce_mean (net, axis=-2, keepdims=True) 58 | net = tf.concat([net_max, net_mean], axis=-1) 59 | if debug: print('Shape {:s} ... Name {:s}'.format(net_max.shape,net_max.name)) 60 | if debug: print('Shape {:s} ... Name {:s}'.format(net_mean.shape,net_mean.name)) 61 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 62 | net = slim.conv2d(inputs = net, 63 | num_outputs = 64, 64 | kernel_size = 1, 65 | stride = 1, 66 | trainable = trainable, 67 | padding = 'VALID', 68 | normalizer_fn = slim.batch_norm, 69 | activation_fn = activation, 70 | scope = 'conv1') 71 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 72 | 73 | return [net_max, net_mean, net] 74 | 75 | def repeat_edge_conv(point_cloud, repeat, k, num_filters, trainable, debug=False): 76 | 77 | repeat = int(repeat) 78 | if not type(k) == type(list()): 79 | k = [int(k)] * repeat 80 | elif not len(k) == repeat: 81 | print('Length of k != repeat') 82 | raise ValueError 83 | if not type(num_filters) == type(list()): 84 | num_filters = [int(num_filters)] * repeat 85 | elif not len(num_filters) == repeat: 86 | print('Length of num_filters != repeat') 87 | raise ValueError 88 | 89 | net = point_cloud 90 | tensors = [] 91 | for i in range(repeat): 92 | scope = 'EdgeConv%d' % i 93 | with tf.variable_scope(scope): 94 | tensors += edge_conv(net, k[i], num_filters[i], trainable, debug=debug) 95 | net = tensors[-1] 96 | net = tf.squeeze(net,axis=-2) 97 | 98 | return tensors 99 | 100 | def repeat_residual_edge_conv(point_cloud, repeat, k, num_filters, trainable, debug=False): 101 | 102 | repeat = int(repeat) 103 | if not type(k) == type(list()): 104 | k = [int(k)] * repeat 105 | elif not len(k) == repeat: 106 | print('Length of k != repeat') 107 | raise ValueError 108 | if not type(num_filters) == type(list()): 109 | num_filters = [int(num_filters)] * repeat 110 | elif not len(num_filters) == repeat: 111 | print('Length of num_filters != repeat') 112 | raise ValueError 113 | 114 | net = point_cloud 115 | tensors = [] 116 | shortcut = None 117 | for i in range(repeat): 118 | scope = 'EdgeConv%d' % i 119 | with tf.variable_scope(scope): 120 | if shortcut is None: 121 | tensors += edge_conv(net, k[i], num_filters[i], trainable, debug=debug) 122 | else: 123 | tensors += edge_conv(net, k[i], num_filters[i], trainable, activation=None, debug=debug) 124 | if not num_filters[i] == num_filters[i-1]: 125 | shortcut = slim.conv2d(inputs = shortcut, 126 | num_outputs = num_filters[i], 127 | kernel_size = 1, 128 | stride = 1, 129 | trainable = trainable, 130 | padding = 'VALID', 131 | normalizer_fn = slim.batch_norm, 132 | activation_fn = None, 133 | scope = 'shortcut') 134 | tensors[-1] = tf.nn.relu(shortcut + tensors[-1]) 135 | 136 | net = tensors[-1] 137 | shortcut = tensors[-1] 138 | net = tf.squeeze(net,axis=-2) 139 | 140 | return tensors 141 | 142 | def fc(net, repeat, num_filters, trainable, debug=False): 143 | 144 | repeat = int(repeat) 145 | if not type(num_filters) == type(list()): 146 | num_filters = [int(num_filters)] * repeat 147 | elif not len(num_filters) == repeat: 148 | print('Length of num_filters != repeat') 149 | raise ValueError 150 | 151 | for i in range(repeat): 152 | scope='FC%d' % i 153 | net = slim.conv2d(inputs = net, 154 | num_outputs = num_filters[i], 155 | kernel_size = 1, 156 | stride = 1, 157 | trainable = trainable, 158 | padding = 'VALID', 159 | normalizer_fn = slim.batch_norm, 160 | scope = scope) 161 | if debug: print('Shape {:s} ... Name {:s}'.format(net.shape,net.name)) 162 | 163 | return net 164 | -------------------------------------------------------------------------------- /dgcnn/trainval.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import dgcnn 5 | import tensorflow as tf 6 | 7 | class trainval(object): 8 | 9 | def __init__(self,flags): 10 | self._flags = flags 11 | 12 | def initialize(self): 13 | 14 | self._ops = {} 15 | 16 | with tf.device('/cpu:0'): 17 | self._optimizer = tf.train.AdamOptimizer(self._flags.LEARNING_RATE) 18 | 19 | self._points_v = [] 20 | self._labels_v = [] 21 | self._weights_v = [] 22 | loss_v = [] 23 | accuracy_v = [] 24 | grad_v = [] 25 | softmax_v = [] 26 | for i, gpu_id in enumerate(self._flags.GPUS): 27 | with tf.device('/gpu:%d' % i): 28 | with tf.name_scope('GPU%d' % gpu_id) as scope: 29 | with tf.variable_scope("dgcnn",reuse=tf.AUTO_REUSE): 30 | points = tf.placeholder(tf.float32, 31 | shape=(self._flags.MINIBATCH_SIZE,None,self._flags.NUM_CHANNEL)) 32 | #shape=(self._flags.MINIBATCH_SIZE,self._flags.NUM_POINT,self._flags.NUM_CHANNEL)) 33 | labels = tf.placeholder(tf.int32, 34 | shape=(self._flags.MINIBATCH_SIZE,None)) 35 | #shape=(self._flags.MINIBATCH_SIZE,self._flags.NUM_POINT)) 36 | self._points_v.append(points) 37 | self._labels_v.append(labels) 38 | pred = dgcnn.model.build(points, self._flags) 39 | softmax = tf.nn.softmax(logits=pred) 40 | softmax_v.append(softmax) 41 | correct = tf.equal(tf.argmax(pred, 2), tf.to_int64(labels)) 42 | accuracy = tf.reduce_mean(tf.cast(correct,tf.float32)) 43 | accuracy_v.append(accuracy) 44 | # If training, compute gradients 45 | if self._flags.TRAIN: 46 | loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=pred, labels=labels) 47 | if self._flags.WEIGHT_KEY: 48 | weight = tf.placeholder(tf.float32, 49 | shape=(self._flags.MINIBATCH_SIZE,None)) 50 | self._weights_v.append(weight) 51 | loss = tf.multiply(loss,weight) 52 | loss = tf.reduce_mean(loss) 53 | loss_v.append(loss) 54 | grad = self._optimizer.compute_gradients(loss) 55 | grad_v.append(grad) 56 | # Softmax 57 | self._softmax_v = softmax_v 58 | # Average loss & accuracy across GPUs 59 | self._loss = tf.add_n(loss_v) / float(len(self._flags.GPUS)) 60 | self._accuracy = tf.add_n(accuracy_v) / float(len(self._flags.GPUS)) 61 | # If training, average gradients across GPUs 62 | if self._flags.TRAIN: 63 | average_grad_v = [] 64 | for grad_and_var_v in zip(*grad_v): 65 | v = [] 66 | for g, _ in grad_and_var_v: 67 | v.append(tf.expand_dims(g,0)) 68 | 69 | grad = tf.reduce_mean(tf.concat(v,0), 0) 70 | 71 | if self._flags.DEBUG: 72 | print('Computing gradients for %s from %d GPUs' % (grad_and_var_v[0][1].name,len(grad_and_var_v))) 73 | average_grad_v.append((grad, grad_and_var_v[0][1])) 74 | 75 | accum_vars = [tf.Variable(v.initialized_value(),trainable=False) for v in tf.trainable_variables()] 76 | self._zero_grad = [v.assign(tf.zeros_like(v)) for v in accum_vars] 77 | self._accum_grad_v = [] 78 | 79 | self._accum_grad_v += [accum_vars[j].assign_add(g[0]) for j,g in enumerate(average_grad_v)] 80 | self._apply_grad = self._optimizer.apply_gradients(zip(accum_vars, tf.trainable_variables())) 81 | 82 | # Merge summary 83 | tf.summary.scalar('accuracy', self._accuracy) 84 | tf.summary.scalar('loss', self._loss) 85 | self._merged_summary=tf.summary.merge_all() 86 | 87 | def feed_dict(self,data,label=None,weight=None): 88 | res = {} 89 | for i,gpu_id in enumerate(self._flags.GPUS): 90 | res[self._points_v [i]] = data [i] 91 | if label is not None: 92 | res[self._labels_v [i]] = label [i] 93 | if weight is not None: 94 | res[self._weights_v [i]] = weight [i] 95 | return res 96 | 97 | def make_summary(self, sess, data, label, weight): 98 | if not self._flags.TRAIN: 99 | raise NotImplementedError 100 | feed_dict = self.feed_dict(data,label,weight) 101 | return sess.run(self._merged_summary,feed_dict=feed_dict) 102 | 103 | def inference(self,sess,data,label=None,weight=None): 104 | feed_dict = self.feed_dict(data,label,weight) 105 | ops = list(self._softmax_v) 106 | if label is not None: 107 | ops += [self._accuracy, self._loss] 108 | return sess.run(ops, feed_dict = feed_dict) 109 | 110 | def accum_gradient(self, sess, data, label, weight=None, summary=False): 111 | if not self._flags.TRAIN: 112 | raise NotImplementedError 113 | 114 | feed_dict = self.feed_dict(data,label,weight) 115 | ops = [self._accum_grad_v] 116 | ops += [self._accuracy, self._loss] 117 | if summary: 118 | ops += [self._merged_summary] 119 | return sess.run(ops, feed_dict = feed_dict) 120 | 121 | def zero_gradients(self, sess): 122 | if not self._flags.TRAIN: 123 | raise NotImplementedError 124 | return sess.run([self._zero_grad]) 125 | 126 | def apply_gradient(self,sess): 127 | if not self._flags.TRAIN: 128 | raise NotImplementedError 129 | return sess.run(self._apply_grad) 130 | -------------------------------------------------------------------------------- /scripts/lsf/continue_dgcnn_vnp.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | MODEL_NAME=$1; # [dgcnn,residual-dgcnn,residual-dgcnn-nofc] 4 | NUM_LAYERS=$2; # recommended: 6 5 | ITERATION=$3; # 6 | WEIGHT_ID=$4; 7 | LEARNING_RATE=$5; 8 | KVALUE=40; 9 | NUM_FILTERS=64; 10 | WEIGHT_KEY=""; # leaving it empty means no weight 11 | 12 | #CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 13 | IN_STORAGE_DIR=/gpfs/slac/staas/fs1/g/neutrino/kterao; 14 | OUT_STORAGE_DIR=/gpfs/slac/staas/fs1/g/neutrino/${USER}; 15 | SWDIR=$HOME/sw 16 | DGCNN_DIR=$SWDIR/dynamic-gcnn; 17 | TRAIN_FILE=dlprod_ppn_v08_p02_train.root; 18 | WEIGHT_FILE=/gpfs/slac/staas/fs1/g/neutrino/kterao/${MODEL_NAME}/res${NUM_LAYERS}/weights/snapshot-${WEIGHT_ID} 19 | WORK_DIR=/scratch/kterao/temp_$$; 20 | DATA_DIR=/scratch/kterao/data_$$; 21 | 22 | source $SWDIR/larcv2/configure.sh; 23 | mkdir -p $WORK_DIR; 24 | mkdir -p $DATA_DIR; 25 | cp $IN_STORAGE_DIR/data/$TRAIN_FILE $DATA_DIR; 26 | cd $WORK_DIR; 27 | echo $CUDA_VISIBLE_DEVICES >> log.txt 28 | echo $DGCNN_DIR/bin/dgcnn.py train -bs 24 -mbs 1 -np -1 --gpus $CUDA_VISIBLE_DEVICES -chks 1000 -chkn 1 -chkh 0.2 -ss 100 -rs 50 -it $ITERATION -ld log -wp weights/snapshot -kv $KVALUE -mn $MODEL_NAME -ecl $NUM_LAYERS -ecf $NUM_FILTERS -io larcv -dkey data -lkey segment -if $DATA_DIR/$TRAIN_FILE -mp $WEIGHT_FILE -lr $LEARNING_RATE >> log.txt; 29 | $DGCNN_DIR/bin/dgcnn.py train -bs 24 -mbs 1 -np -1 --gpus $CUDA_VISIBLE_DEVICES -chks 1000 -chkn 1 -chkh 0.2 -ss 100 -rs 50 -it $ITERATION -ld log -wp weights/snapshot -kv $KVALUE -mn $MODEL_NAME -ecl $NUM_LAYERS -ecf $NUM_FILTERS -io larcv -dkey data -lkey segment -if $DATA_DIR/$TRAIN_FILE -mp $WEIGHT_FILE -lr $LEARNING_RATE >> log.txt; 30 | cd ..; 31 | mkdir -p $OUT_STORAGE_DIR; 32 | cp -r $WORK_DIR $OUT_STORAGE_DIR; 33 | if [ $? -eq 1 ]; then 34 | cp -r $WORK_DIR $HOME; 35 | fi 36 | if [ $? -eq 0 ]; then 37 | rm -rf $WORK_DIR; 38 | fi 39 | rm -rf $DATA_DIR; 40 | 41 | -------------------------------------------------------------------------------- /scripts/lsf/inference_dgcnn.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | MODEL_NAME=$1; # [dgcnn,residual-dgcnn,residual-dgcnn-nofc] 4 | NUM_LAYERS=$2; # recommended: 6 5 | ITERATION=$3; # 6 | WEIGHT_ID=$4; 7 | LEARNING_RATE=$5; 8 | KVALUE=40; 9 | NUM_FILTERS=64; 10 | WEIGHT_KEY=""; # leaving it empty means no weight 11 | GPUS=`python -c "import os;print len(os.environ['CUDA_VISIBLE_DEVICES'].split(','))"`; 12 | 13 | IN_STORAGE_DIR=/gpfs/slac/staas/fs1/g/neutrino/kterao; 14 | OUT_STORAGE_DIR=/gpfs/slac/staas/fs1/g/neutrino/${USER}; 15 | SWDIR=$HOME/sw; 16 | DGCNN_DIR=$SWDIR/dynamic-gcnn; 17 | WEIGHT_FILE=/gpfs/slac/staas/fs1/g/neutrino/kterao/${MODEL_NAME}/res${NUM_LAYERS}/weights/snapshot-${WEIGHT_ID} 18 | TEST_FILE=dlprod_ppn_v08_p02_test.root; 19 | WORK_DIR=/scratch/kterao/temp_$$; 20 | DATA_DIR=/scratch/kterao/data_$$; 21 | 22 | 23 | source $SWDIR/larcv2/configure.sh; 24 | mkdir -p $WORK_DIR; 25 | mkdir -p $DATA_DIR; 26 | cp $IN_STORAGE_DIR/data/$TEST_FILE $DATA_DIR; 27 | cd $WORK_DIR; 28 | echo $CUDA_VISIBLE_DEVICES >> log.txt; 29 | echo $GPUS >> log.txt; 30 | echo $DGCNN_DIR/bin/dgcnn.py inference -bs $GPUS -mbs 1 -np -1 --gpus $CUDA_VISIBLE_DEVICES -rs 1 -it $ITERATION -ld log -kv $KVALUE -ecl $NUM_LAYERS -ecf $NUM_FILTERS -io larcv -dkey data -lkey segment -if $DATA_DIR/$TEST_FILE -mp $WEIGHT_FILE -mn $MODEL_NAME -sh 0 >> log.txt; 31 | $DGCNN_DIR/bin/dgcnn.py inference -bs $GPUS -mbs 1 -np -1 --gpus $CUDA_VISIBLE_DEVICES -rs 1 -it $ITERATION -ld log -kv $KVALUE -ecl $NUM_LAYERS -ecf $NUM_FILTERS -io larcv -dkey data -lkey segment -if $DATA_DIR/$TEST_FILE -mp $WEIGHT_FILE -mn $MODEL_NAME -sh 0 >> log.txt; 32 | cd ..; 33 | mkdir -p $OUT_STORAGE_DIR; 34 | cp -r $WORK_DIR $OUT_STORAGE_DIR; 35 | if [ $? -eq 1 ]; then 36 | cp -r $WORK_DIR $HOME; 37 | fi 38 | if [ $? -eq 0 ]; then 39 | rm -rf $WORK_DIR; 40 | fi 41 | rm -rf $DATA_DIR; 42 | -------------------------------------------------------------------------------- /scripts/lsf/train_dgcnn.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | MODEL_NAME=$1; # [dgcnn,residual-dgcnn,residual-dgcnn-nofc] 4 | NUM_LAYERS=$2; # recommended: 6 5 | ITERATION=$3; # 6 | WEIGHT_ID=$4; 7 | LEARNING_RATE=$5; 8 | KVALUE=40; 9 | NUM_FILTERS=64; 10 | WEIGHT_KEY=""; # leaving it empty means no weight 11 | 12 | #CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 13 | IN_STORAGE_DIR=/gpfs/slac/staas/fs1/g/neutrino/kterao; 14 | OUT_STORAGE_DIR=/gpfs/slac/staas/fs1/g/neutrino/${USER}; 15 | SWDIR=$HOME/sw 16 | DGCNN_DIR=$SWDIR/dynamic-gcnn; 17 | TRAIN_FILE=dlprod_ppn_v08_p02_train_slim.root; 18 | WORK_DIR=/scratch/kterao/temp_$$; 19 | DATA_DIR=/scratch/kterao/data_$$; 20 | 21 | source $SWDIR/larcv2/configure.sh; 22 | mkdir -p $WORK_DIR; 23 | mkdir -p $DATA_DIR; 24 | cp $IN_STORAGE_DIR/data/$TRAIN_FILE $DATA_DIR; 25 | cd $WORK_DIR; 26 | echo $CUDA_VISIBLE_DEVICES >> log.txt 27 | echo $DGCNN_DIR/bin/dgcnn.py train -bs 24 -mbs 1 -np -1 --gpus $CUDA_VISIBLE_DEVICES -chks 1000 -chkn 1 -chkh 0.2 -ss 100 -rs 50 -it $ITERATION -ld log -wp weights/snapshot -kv $KVALUE -mn $MODEL_NAME -ecl $NUM_LAYERS -ecf $NUM_FILTERS -io larcv -dkey data -lkey segment -wkey $WEIGHT_KEY -if $DATA_DIR/$TRAIN_FILE >> log.txt; 28 | $DGCNN_DIR/bin/dgcnn.py train -bs 24 -mbs 1 -np -1 --gpus $CUDA_VISIBLE_DEVICES -chks 1000 -chkn 1 -chkh 0.2 -ss 100 -rs 50 -it $ITERATION -ld log -wp weights/snapshot -kv $KVALUE -mn $MODEL_NAME -ecl $NUM_LAYERS -ecf $NUM_FILTERS -io larcv -dkey data -lkey segment -wkey $WEIGHT_KEY -if $DATA_DIR/$TRAIN_FILE >> log.txt; 29 | cd ..; 30 | mkdir -p $OUT_STORAGE_DIR; 31 | cp -r $WORK_DIR $OUT_STORAGE_DIR; 32 | if [ $? -eq 1 ]; then 33 | cp -r $WORK_DIR $HOME; 34 | fi 35 | if [ $? -eq 0 ]; then 36 | rm -rf $WORK_DIR; 37 | fi 38 | rm -rf $DATA_DIR; 39 | 40 | -------------------------------------------------------------------------------- /scripts/plot_csv.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import pandas as pd 3 | df = pd.read_csv(sys.argv[1]) 4 | 5 | print df.describe() 6 | --------------------------------------------------------------------------------