cellmap_flow.inference.runner

Running a model on a region: the device, the warmup and the forward.

ModelRunner puts a ModelConfig’s model on the device, reads its output activation off its network’s graph (ModelConfig.output_graph, portable.activation), warms it up (the same forward checks the declared shapes and probes the output’s range, which says the activation when the graph does not), and predicts an output region from a RawInput (cellmap_flow.io.raw; in the server, its BlockInput): the model’s raw output, (C, *spatial), with no postprocessing and no process-wide state. cellmap_flow.inferencer.Inferencer adds the normalization and postprocessing chain on top.

DeviceSlots is how many chunks may use the device at once, and in what order; the inference server passes one, blockwise workers none.

Importing this module is cheap: torch is imported when it is first needed.

Attributes

logger
GPU_SLOTS_ENV
HALF_PRECISION_TOLERANCE

Classes

DeviceSlots	Lets at most n chunks use the device at once, in the order they asked.
ModelRunner

Functions

predict(read_roi, write_roi, config, **kwargs)

The default forward: read read_roi from kwargs["idi"], run

Module Contents

cellmap_flow.inference.runner.logger

cellmap_flow.inference.runner.GPU_SLOTS_ENV = 'CELLMAP_FLOW_GPU_SLOTS'

cellmap_flow.inference.runner.HALF_PRECISION_TOLERANCE = 0.01

class cellmap_flow.inference.runner.DeviceSlots(n=1)

Lets at most n chunks use the device at once, in the order they asked.

The inference server answers every request on its own thread, and neuroglancer asks for 6 chunks at once over HTTP/1.1, or dozens through an HTTP/2 proxy. Unbounded, k concurrent forwards interleave on the GPU and finish together, after about k forwards’ time, so the first chunks of a view all appeared at once, late. Each forward also holds its own activations, and a dozen at once ran an 11 GB card out of memory. One slot makes the GPU compute chunks one after another, the first one asked for first, at the same total throughput: reading and normalizing the next chunk, which stay outside the slot, overlap the current forward.

Strictly first come, first served: a thread that arrives while a slot is free still waits behind anyone already waiting. A plain Semaphore lets newcomers barge ahead.

The server takes n from CELLMAP_FLOW_GPU_SLOTS (default 1) when it starts. More than one slot only helps a model too small to keep the device busy on its own.

Neuroglancer drops the requests for chunks that have left the view after a pan. Dropping those chunks before they reach a slot is the server’s queue’s job (chunkmirage’s), which admits chunks to the ops in the order they were asked for and drops those nobody waits for any more; a chunk already on the device finishes.

n = 1

classmethod from_env()

hold()

cellmap_flow.inference.runner.predict(read_roi, write_roi, config, **kwargs)

The default forward: read read_roi from kwargs["idi"], run config.model on kwargs["device"], return (C, *spatial).

idi is what the model reads from: a RawInput (cellmap_flow.io.raw), or the server’s BlockInput, anything with to_ndarray_ts(roi); the keyword keeps its old name, which scripts read. A script may define its own predict with this signature instead; it is then given idi, device and use_half_prediction.

class cellmap_flow.inference.runner.ModelRunner(model_config, use_half_prediction=False, device_slots=None, half_precision=False)

device_slots = None

device

use_half_prediction = False

model_config

output_range = None

output_class = None

output_class_from = None

output_op = None

half_precision

autocast_dtype = None

context

optimize_model()

predict(idi, out_roi)

The model’s raw output for out_roi: (C, *spatial), not postprocessed.

idi (a RawInput, or anything with to_ndarray_ts(roi)) is read around out_roi with the model’s context, through whatever input chain it reads with.

process_chunk_basic(idi, roi)