def eval_model_zip(model_zip: ZipFile):
with TemporaryDirectory() as tempdir:
temp_path = Path(tempdir)
model_zip.extractall(temp_path)
spec_file_str = guess_model_path(
[str(file_name) for file_name in temp_path.glob("*")])
bioimageio_model = load_model(spec_file_str)
return get_nn_instance(bioimageio_model)
def __init__(self, *, bioimageio_model: nodes.Model, devices=Sequence[str]):
self._internal_output_axes = bioimageio_model.outputs[0].axes
spec = bioimageio_model
self.model = get_nn_instance(bioimageio_model)
self.devices = [torch.device(d) for d in devices]
self.model.to(self.devices[0])
assert isinstance(self.model, torch.nn.Module)
weights = spec.weights.get("pytorch_state_dict")
if weights is not None and weights.source:
state = torch.load(weights.source, map_location=self.devices[0])
self.model.load_state_dict(state)
def __init__(self, *, bioimageio_model: nodes.Model, devices=List[str]):
spec = bioimageio_model
self.name = spec.name
spec.inputs[0]
_output = spec.outputs[0]
# FIXME: TF probably uses different axis names
self._internal_output_axes = _output.axes
self.model = get_nn_instance(bioimageio_model)
self.devices = []
tf_model = tf.keras.models.load_model(
spec.weights["tensorflow_saved_model_bundle"].source)
self.model.set_model(tf_model)
def convert_weights_to_onnx(model_yaml: Union[str, Path],
output_path: Union[str, Path],
opset_version: Union[str, None] = 12,
use_tracing: bool = True,
verbose: bool = True):
""" Convert model weights from format 'pytorch_state_dict' to 'onnx'.
Arguments:
model_yaml: location of the model.yaml file with bioimage.io spec
output_path: where to save the onnx weights
opset_version: onnx opset version
use_tracing: whether to use tracing or scripting to export the onnx format
verbose: be verbose during the onnx export
"""
spec = load_model(model_yaml)
with torch.no_grad():
# load input and expected output data
input_data = np.load(spec.test_inputs[0]).astype('float32')
input_tensor = torch.from_numpy(input_data)
# instantiate and generate the expected output
model = get_nn_instance(spec)
state = torch.load(spec.weights['pytorch_state_dict'].source)
model.load_state_dict(state)
expected_output = model(input_tensor).numpy()
if use_tracing:
torch.onnx.export(model,
input_tensor,
output_path,
verbose=verbose,
opset_version=opset_version)
else:
raise NotImplementedError
# check the onnx model
sess = rt.InferenceSession(output_path)
input_name = sess.get_inputs()[0].name
output = sess.run(None, {input_name: input_data})[0]
try:
assert_array_almost_equal(expected_output, output, decimal=4)
return 0
except AssertionError as e:
msg = f"The onnx weights were exported, but results before and after conversion do not agree:\n {str(e)}"
warnings.warn(msg)
return 1
def convert_weights_to_torchscript(
model_yaml: Union[str, Path],
output_path: Union[str, Path],
use_tracing: bool = True
):
""" Convert model weights from format 'pytorch_state_dict' to 'torchscript'.
"""
spec = load_model(model_yaml)
with torch.no_grad():
# load input and expected output data
input_data = np.load(spec.test_inputs[0]).astype('float32')
input_data = torch.from_numpy(input_data)
# instantiate model and get reference output
model = get_nn_instance(spec)
state = torch.load(spec.weights['pytorch_state_dict'].source)
model.load_state_dict(state)
# get the expected output to validate the torchscript weights
expected_output = model(input_data)
# make scripted model
if use_tracing:
scripted_model = torch.jit.trace(model, input_data)
else:
scripted_model = torch.jit.script(model)
# check the scripted model
output = scripted_model(input_data).numpy()
# save the torchscript model
scripted_model.save(output_path)
try:
assert_array_almost_equal(expected_output, output, decimal=4)
return 0
except AssertionError as e:
msg = f"The onnx weights were exported, but results before and after conversion do not agree:\n {str(e)}"
warnings.warn(msg)
return 1
def simple_training(
bioimageio_model: ModelWithKwargs, n_iterations: int, batch_size: int, num_workers: int, out_file: Union[str, Path, IO[bytes]]
) -> torch.nn.Module:
""" Simplified training loop.
"""
if isinstance(out_file, str) or isinstance(out_file, Path):
out_file = Path(out_file)
out_file.parent.mkdir(exist_ok=True)
model = get_nn_instance(bioimageio_model)
# instantiate all training parameters from the training config
setup = bioimageio_model.spec.training.setup
sampler = get_nn_instance(setup.sampler)
preprocess = [get_nn_instance(prep) for prep in setup.preprocess]
postprocess = [get_nn_instance(post) for post in setup.postprocess]
losses = [get_nn_instance(loss_prep) for loss_prep in setup.losses]
optimizer = get_nn_instance(setup.optimizer, params=model.parameters())
# build the data-loader from our sampler
loader = DataLoader(sampler, shuffle=True, num_workers=num_workers, batch_size=batch_size)
# run the training loop
for ii in trange(n_iterations):
x, y = next(iter(loader))
optimizer.zero_grad()
x, y = apply_transformations(preprocess, x, y)
out = model(x)
out, y = apply_transformations(postprocess, out, y)
losses = apply_transformations(losses, out, y)
ll = sum(losses)
ll.backward()
optimizer.step()
# save model weights
torch.save(model.state_dict(), out_file)
return model