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())
class CaffePredict(object):
'''Augments a batch with network predictions.
Args:
prototxt (string): Filename of the network prototxt.
weights (string): Filename of the network weights.
input_key (string): Name of the input layer of the network.
output_key (string): Name of the output layer of the network.
gpu (int): Which GPU to use.
'''
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])
def __call__(self, input_data):
assert isinstance(input_data, np.ndarray)
self.net_io.set_inputs({self.input_key: input_data})
self.net.forward()
output = self.net_io.get_outputs()[self.output_key]
assert isinstance(output, np.ndarray)
# remove batch-dimension
if output.ndim == 5:
output = output[0]
assert output.ndim == 4
return output.astype('float32')
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 start(self):
logger.info("Initializing solver...")
if self.use_gpu is not None:
logger.debug("Train 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.solver = caffe.get_solver(self.solver_parameters)
if self.solver_parameters.resume_from is not None:
logger.debug("Train process: restoring solver state from %s",
self.solver_parameters.resume_from)
self.solver.restore(self.solver_parameters.resume_from)
names_net_outputs = self.outputs.keys() + self.gradients.keys()
self.net_io = NetIoWrapper(self.solver.net, names_net_outputs)
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])
def __train(self, use_gpu):
start = time.time()
if not self.solver_initialized:
logger.info("Initializing solver...")
if use_gpu is not None:
logger.debug("Train 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.solver = caffe.get_solver(self.solver_parameters)
if self.solver_parameters.resume_from is not None:
logger.debug("Train process: restoring solver state from " +
self.solver_parameters.resume_from)
self.solver.restore(self.solver_parameters.resume_from)
self.net_io = NetIoWrapper(self.solver.net)
self.solver_initialized = True
batch, request = self.batch_in.get()
data = {
'data':
batch.volumes[VolumeType.RAW].data[np.newaxis, np.newaxis, :],
'aff_label':
batch.volumes[VolumeType.GT_AFFINITIES].data[np.newaxis, :],
}
if self.solver_parameters.train_state.get_stage(0) == 'euclid':
logger.debug(
"Train process: preparing input data for Euclidean training")
self.__prepare_euclidean(batch, data)
else:
logger.debug(
"Train process: preparing input data for Malis training")
self.__prepare_malis(batch, data)
self.net_io.set_inputs(data)
loss = self.solver.step(1)
# self.__consistency_check()
output = self.net_io.get_outputs()
batch.volumes[VolumeType.PRED_AFFINITIES] = Volume(
output['aff_pred'][0],
batch.volumes[VolumeType.GT_AFFINITIES].roi,
batch.volumes[VolumeType.GT_AFFINITIES].resolution,
interpolate=True)
batch.loss = loss
batch.iteration = self.solver.iter
if VolumeType.LOSS_GRADIENT in request.volumes:
diffs = self.net_io.get_output_diffs()
batch.volumes[VolumeType.LOSS_GRADIENT] = Volume(
diffs['aff_pred'][0],
batch.volumes[VolumeType.GT_AFFINITIES].roi,
batch.volumes[VolumeType.GT_AFFINITIES].resolution,
interpolate=True)
time_of_iteration = time.time() - start
logger.info("Train process: iteration=%d loss=%f time=%f" %
(self.solver.iter, batch.loss, time_of_iteration))
return batch
class Train(BatchFilter):
'''Performs one training iteration for each batch that passes through.
Adds the predicted affinities to the batch.
'''
def __init__(self, solver_parameters, use_gpu=None):
# start training as a producer pool, so that we can gracefully exit if
# anything goes wrong
self.worker = ProducerPool([lambda gpu=use_gpu: self.__train(gpu)],
queue_size=1)
self.batch_in = multiprocessing.Queue(maxsize=1)
self.solver_parameters = solver_parameters
self.solver_initialized = False
def setup(self):
self.worker.start()
def teardown(self):
self.worker.stop()
def prepare(self, request):
# remove request parts that we provide
for volume_type in [
VolumeType.LOSS_GRADIENT, VolumeType.PRED_AFFINITIES
]:
if volume_type in request.volumes:
del request.volumes[volume_type]
def process(self, batch, request):
self.batch_in.put((batch, request))
try:
out = self.worker.get()
except WorkersDied:
raise TrainProcessDied()
batch.volumes[VolumeType.PRED_AFFINITIES] = out.volumes[
VolumeType.PRED_AFFINITIES]
if VolumeType.LOSS_GRADIENT in request.volumes:
batch.volumes[VolumeType.LOSS_GRADIENT] = out.volumes[
VolumeType.LOSS_GRADIENT]
batch.loss = out.loss
batch.iteration = out.iteration
def __train(self, use_gpu):
start = time.time()
if not self.solver_initialized:
logger.info("Initializing solver...")
if use_gpu is not None:
logger.debug("Train 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.solver = caffe.get_solver(self.solver_parameters)
if self.solver_parameters.resume_from is not None:
logger.debug("Train process: restoring solver state from " +
self.solver_parameters.resume_from)
self.solver.restore(self.solver_parameters.resume_from)
self.net_io = NetIoWrapper(self.solver.net)
self.solver_initialized = True
batch, request = self.batch_in.get()
data = {
'data':
batch.volumes[VolumeType.RAW].data[np.newaxis, np.newaxis, :],
'aff_label':
batch.volumes[VolumeType.GT_AFFINITIES].data[np.newaxis, :],
}
if self.solver_parameters.train_state.get_stage(0) == 'euclid':
logger.debug(
"Train process: preparing input data for Euclidean training")
self.__prepare_euclidean(batch, data)
else:
logger.debug(
"Train process: preparing input data for Malis training")
self.__prepare_malis(batch, data)
self.net_io.set_inputs(data)
loss = self.solver.step(1)
# self.__consistency_check()
output = self.net_io.get_outputs()
batch.volumes[VolumeType.PRED_AFFINITIES] = Volume(
output['aff_pred'][0],
batch.volumes[VolumeType.GT_AFFINITIES].roi,
batch.volumes[VolumeType.GT_AFFINITIES].resolution,
interpolate=True)
batch.loss = loss
batch.iteration = self.solver.iter
if VolumeType.LOSS_GRADIENT in request.volumes:
diffs = self.net_io.get_output_diffs()
batch.volumes[VolumeType.LOSS_GRADIENT] = Volume(
diffs['aff_pred'][0],
batch.volumes[VolumeType.GT_AFFINITIES].roi,
batch.volumes[VolumeType.GT_AFFINITIES].resolution,
interpolate=True)
time_of_iteration = time.time() - start
logger.info("Train process: iteration=%d loss=%f time=%f" %
(self.solver.iter, batch.loss, time_of_iteration))
return batch
def __prepare_euclidean(self, batch, data):
gt_affinities = batch.volumes[VolumeType.GT_AFFINITIES]
# initialize error scale with 1s
error_scale = np.ones(gt_affinities.data.shape, dtype=np.float)
# set error_scale to 0 in masked-out areas
if VolumeType.GT_MASK in batch.volumes:
self.__mask_error_scale(error_scale,
batch.volumes[VolumeType.GT_MASK].data)
if VolumeType.GT_IGNORE in batch.volumes:
self.__mask_error_scale(error_scale,
batch.volumes[VolumeType.GT_IGNORE].data)
# in the masked-in area, compute the fraction of positive samples
masked_in = error_scale.sum()
num_pos = (gt_affinities.data * error_scale).sum()
frac_pos = float(num_pos) / masked_in if masked_in > 0 else 0
frac_pos = np.clip(frac_pos, 0.05, 0.95)
frac_neg = 1.0 - frac_pos
# compute the class weights for positive and negative samples
w_pos = 1.0 / (2.0 * frac_pos)
w_neg = 1.0 / (2.0 * frac_neg)
# scale the masked-in error_scale with the class weights
error_scale *= (data >= 0.5) * w_pos + (data < 0.5) * w_neg
data['scale'] = error_scale[np.newaxis, :]
def __mask_error_scale(self, error_scale, mask):
for d in range(error_scale.shape[0]):
error_scale[d] *= mask
def __prepare_malis(self, batch, data):
gt_labels = batch.volumes[VolumeType.GT_LABELS]
next_id = gt_labels.data.max() + 1
gt_pos_pass = gt_labels.data
if VolumeType.GT_IGNORE in batch.volumes:
gt_neg_pass = np.array(gt_labels.data)
gt_neg_pass[batch.volumes[VolumeType.GT_IGNORE].data ==
0] = next_id
else:
gt_neg_pass = gt_pos_pass
data['comp_label'] = np.array([[gt_neg_pass, gt_pos_pass]])
data['nhood'] = batch.affinity_neighborhood[np.newaxis, np.newaxis, :]
# Why don't we update gt_affinities in the same way?
# -> not needed
#
# GT affinities are all 0 in the masked area (because masked area is
# assumed to be set to background in batch.gt).
#
# In the negative pass:
#
# We set all affinities inside GT regions to 1. Affinities in masked
# area as predicted. Belongs to one forground region (introduced
# above). But we only count loss on edges connecting different labels
# -> loss in masked-out area only from outside regions.
#
# In the positive pass:
#
# We set all affinities outside GT regions to 0 -> no loss in masked
# out area.
def __consistency_check(self):
diffs = self.net_io.get_outputs()
for k in diffs:
assert not np.isnan(
diffs[k]).any(), "Detected NaN in output diff " + k
class Predict(BatchFilter):
'''Augments the batch with the predicted affinities.
'''
def __init__(self, prototxt, weights, use_gpu=None):
for f in [prototxt, weights]:
if not os.path.isfile(f):
raise RuntimeError("%s does not exist" % f)
# start prediction as a producer pool, so that we can gracefully exit if
# anything goes wrong
self.worker = ProducerPool([lambda gpu=use_gpu: self.__predict(gpu)],
queue_size=1)
self.batch_in = multiprocessing.Queue(maxsize=1)
self.prototxt = prototxt
self.weights = weights
self.net_initialized = False
def setup(self):
self.worker.start()
def teardown(self):
self.worker.stop()
def prepare(self, request):
# remove request parts that we provide
if VolumeTypes.PRED_AFFINITIES in request.volumes:
del request.volumes[VolumeTypes.PRED_AFFINITIES]
def process(self, batch, request):
self.batch_in.put(batch)
try:
out = self.worker.get()
except WorkersDied:
raise PredictProcessDied()
affs = out.volumes[VolumeTypes.PRED_AFFINITIES]
affs.roi = request.volumes[VolumeTypes.PRED_AFFINITIES]
affs.resolution = batch.volumes[VolumeTypes.RAW].resolution
batch.volumes[VolumeTypes.PRED_AFFINITIES] = affs
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
class Train(GenericTrain):
'''Caffe implementation of :class:`gunpowder.nodes.Train`.
Args:
solver_parameters (:class:``SolverParameters``): Parameters of the
solver to use for training, contains the network description as
well.
inputs (dict): Dictionary from names of input layers in the network to
:class:``ArrayKey`` or batch attribute name as string.
outputs (dict): Dictionary from the names of output layers in the
network to :class:``ArrayKey``. New arrays will be generated by
this node for each entry (if requested downstream).
gradients (dict): Dictionary from the names of output layers in the
network to :class:``ArrayKey``. New arrays containing the
gradient of an output with respect to the loss will be generated by
this node for each entry (if requested downstream).
array_specs (dict, optional): An optional dictionary of
:class:`ArrayKey` to :class:`ArraySpec` to set the array specs
generated arrays (``outputs`` and ``gradients``). This is useful
to set the ``voxel_size``, for example, if they differ from the
voxel size of the input arrays. Only fields that are not ``None``
in the given :class:`ArraySpec` will be used.
use_gpu (int): Which GPU to use. Set to ``None`` for CPU mode.
'''
def __init__(self,
solver_parameters,
inputs,
outputs,
gradients,
array_specs=None,
use_gpu=None):
super(Train, self).__init__(inputs,
outputs,
gradients,
array_specs,
spawn_subprocess=True)
self.solver_parameters = solver_parameters
self.use_gpu = use_gpu
self.solver = None
self.net_io = None
def start(self):
logger.info("Initializing solver...")
if self.use_gpu is not None:
logger.debug("Train 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.solver = caffe.get_solver(self.solver_parameters)
if self.solver_parameters.resume_from is not None:
logger.debug("Train process: restoring solver state from %s",
self.solver_parameters.resume_from)
self.solver.restore(self.solver_parameters.resume_from)
names_net_outputs = self.outputs.keys() + self.gradients.keys()
self.net_io = NetIoWrapper(self.solver.net, names_net_outputs)
def train_step(self, batch, request):
data = {}
for input_name, network_input in self.inputs.items():
if isinstance(network_input, ArrayKey):
if network_input in batch.arrays:
data[input_name] = batch.arrays[network_input].data
else:
logger.warn(
"batch does not contain %s, input %s will not "
"be set", network_input, input_name)
elif isinstance(network_input, np.ndarray):
data[input_name] = network_input
elif isinstance(network_input, str):
data[input_name] = getattr(batch, network_input)
else:
raise Exception(
"Unknown network input type {}, can't be given to "
"network".format(network_input))
self.net_io.set_inputs(data)
loss = self.solver.step(1)
# self.__consistency_check()
requested_outputs = {
name: array_key
for name, array_key in self.outputs.items()
if array_key in request.array_specs
}
if requested_outputs:
output = self.net_io.get_outputs()
for output_name, array_key in requested_outputs.items():
spec = self.spec[array_key].copy()
spec.roi = request[array_key].roi
batch.arrays[array_key] = Array(
output[output_name][0], # strip #batch dimension
spec)
requested_gradients = {
name: array_key
for name, array_key in self.gradients.items()
if array_key in request.array_specs
}
if requested_gradients:
diffs = self.net_io.get_output_diffs()
for output_name, array_key in requested_gradients.items():
spec = self.spec[array_key].copy()
spec.roi = request[array_key].roi
batch.arrays[array_key] = Array(
diffs[output_name][0], # strip #batch dimension
spec)
batch.loss = loss
batch.iteration = self.solver.iter
def __consistency_check(self):
diffs = self.net_io.get_output_diffs()
for k in diffs:
assert not np.isnan(
diffs[k]).any(), "Detected NaN in output diff " + k
class StupidPredict(object):
'''Augments a batch with network predictions.
Args:
prototxt (string): Filename of the network prototxt.
weights (string): Filename of the network weights.
inputs (dict): Dictionary from the names of input layers in the
network to :class:``VolumeType`` or batch attribute name as string.
outputs (dict): Dictionary from the names of output layers in the
network to :class:``VolumeType``. New volumes will be generated by
this node for each entry (if requested downstream).
volume_specs (dict, optional): An optional dictionary of
:class:`VolumeType` to :class:`VolumeSpec` to set the volume specs
generated volumes (``outputs``). This is useful to set the
``voxel_size``, for example, if they differ from the voxel size of
the input volumes. Only fields that are not ``None`` in the given
:class:`VolumeSpec` will be used.
use_gpu (int): Which GPU to use. Set to ``None`` for CPU mode.
'''
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 __call__(self, input_data):
assert isinstance(input_data, dict)
self.net_io.set_inputs(
{input_name: data
for input_name, data in input_data.items()})
self.net.forward()
output = self.net_io.get_outputs()
return output
class Predict(GenericPredict):
'''Augments a batch with network predictions.
Args:
prototxt (string): Filename of the network prototxt.
weights (string): Filename of the network weights.
inputs (dict): Dictionary from the names of input layers in the
network to :class:``ArrayKey`` or batch attribute name as string.
outputs (dict): Dictionary from the names of output layers in the
network to :class:``ArrayKey``. New arrays will be generated by
this node for each entry (if requested downstream).
array_specs (dict, optional): An optional dictionary of
:class:`ArrayKey` to :class:`ArraySpec` to set the array specs
generated arrays (``outputs``). This is useful to set the
``voxel_size``, for example, if they differ from the voxel size of
the input arrays. Only fields that are not ``None`` in the given
:class:`ArraySpec` will be used.
use_gpu (int): Which GPU to use. Set to ``None`` for CPU mode.
'''
def __init__(self,
prototxt,
weights,
inputs,
outputs,
array_specs=None,
use_gpu=None):
super(Predict, self).__init__(inputs,
outputs,
array_specs,
spawn_subprocess=True)
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
self.use_gpu = use_gpu
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, batch, request):
self.net_io.set_inputs({
input_name: batch.arrays[array_key].data
for input_name, array_key in self.inputs.items()
})
self.net.forward()
output = self.net_io.get_outputs()
for output_name, array_key in self.outputs.items():
spec = self.spec[array_key].copy()
spec.roi = request[array_key].roi
batch.arrays[array_key] = Array(
output[output_name][0], # strip #batch dimension
spec)
return batch