def __init__(self, prototxt, weights, inputs, outputs, volume_specs=None, use_gpu=None): for f in [prototxt, weights]: if not os.path.isfile(f): raise RuntimeError("%s does not exist" % f) self.prototxt = prototxt self.weights = weights self.inputs = inputs self.outputs = outputs if use_gpu is not None: logger.debug("Predict process: using GPU %d" % use_gpu) caffe.enumerate_devices(False) caffe.set_devices((use_gpu, )) caffe.set_mode_gpu() caffe.select_device(use_gpu, False) self.net = caffe.Net(self.prototxt, self.weights, caffe.TEST) self.net_io = NetIoWrapper(self.net, self.outputs.values())
def start(self): logger.info("Initializing solver...") if self.use_gpu is not None: logger.debug("Predict process: using GPU %d" % self.use_gpu) caffe.enumerate_devices(False) caffe.set_devices((self.use_gpu, )) caffe.set_mode_gpu() caffe.select_device(self.use_gpu, False) self.net = caffe.Net(self.prototxt, self.weights, caffe.TEST) self.net_io = NetIoWrapper(self.net, self.outputs.keys())
def __predict(self, use_gpu): if not self.net_initialized: logger.info("Initializing solver...") if use_gpu is not None: logger.debug("Predict process: using GPU %d" % use_gpu) caffe.enumerate_devices(False) caffe.set_devices((use_gpu, )) caffe.set_mode_gpu() caffe.select_device(use_gpu, False) self.net = caffe.Net(self.prototxt, self.weights, caffe.TEST) self.net_io = NetIoWrapper(self.net) self.net_initialized = True start = time.time() batch = self.batch_in.get() fetch_time = time.time() - start self.net_io.set_inputs({ 'data': batch.volumes[VolumeTypes.RAW].data[np.newaxis, np.newaxis, :], }) self.net.forward() output = self.net_io.get_outputs() predict_time = time.time() - start logger.info( "Predict process: time=%f (including %f waiting for batch)" % (predict_time, fetch_time)) assert len( output['aff_pred'].shape ) == 5, "Got affinity prediction with unexpected number of dimensions, should be 1 (direction) + 3 (spatial) + 1 (batch, not used), but is %d" % len( output['aff_pred'].shape) batch.volumes[VolumeTypes.PRED_AFFINITIES] = Volume( output['aff_pred'][0], Roi(), (1, 1, 1)) return batch
def __init__(self, prototxt, weights, input_key, output_key, gpu): # TODO validate that gpu is available assert os.path.exists(prototxt) assert os.path.exists(weights) for f in [prototxt, weights]: if not os.path.isfile(f): raise RuntimeError("%s does not exist" % f) self.prototxt = prototxt self.weights = weights self.input_key = input_key self.output_key = output_key caffe.enumerate_devices(False) caffe.set_devices((gpu, )) caffe.set_mode_gpu() caffe.select_device(gpu, False) self.net = caffe.Net(self.prototxt, self.weights, caffe.TEST) self.net_io = NetIoWrapper(self.net, [self.output_key])