Convert to float16 model

[Bludator]

I have a pytorch model that I would like to convert to tflite model with float16 weights and activations.

I am actually not sure if this is a bug report or feature request. There is more or less no documentation on this matter. So I don't know if I am doing something wrong, there is bug or it is simply not there 🤔 .

I have been testing it on dummy model:

Model code

class SimpleCNN(nn.Module):
	def __init__(self):
		super(SimpleCNN, self).__init__()
		self.model = nn.Sequential(
			nn.Conv2d(3, 1, kernel_size=8, stride=8, padding=2),
			nn.ReLU(),
			nn.Flatten(),
			nn.Linear(32 * 32, 10),
		)

	def forward(self, x):
		return self.model(x)

I tried to set the model in pytorch to float16 first i.e.

edge_model = ai_edge_torch.convert(
	model.half(), tuple([x.half() for x in sample_inputs])
)

which fails catastrophically with

loc("main.SimpleCNN/torch.nn.modules.container.Sequential_model/torch.nn.modules.conv.Conv2d_0;"("conv2d"("SimpleCNN/main.py":17:0))): error: failed to legalize operation 'tfl.transpose' that was explicitly marked illegal

Then I have tried to persuade the TensorFlow Lite Quantization to do the thing

converter_flags["optimizations"] = [tf.lite.Optimize.DEFAULT]
converter_flags["target_spec.supported_types"] = [tf.float16]

edge_model = ai_edge_torch.convert(
	model, sample_inputs, _ai_edge_converter_flags=converter_flags
)

This does not crashes but the converted model does not have float16 weights it has mix of float32 and int8 weight types.

[aravindmurali22]

Hi @Bludator,

I have successfully reproduced the issue on my end. Here is what I found:

1. model.half(): It appears this fails because the current pipeline is optimized for float32 input models.
2. Flags: I confirmed that using dot notation (e.g., "target_spec.supported_types") causes the flag to be ignored. This happens because the system interprets it as a literal attribute name rather than traversing the object structure.

Could you please try using a nested dictionary instead? Here is an example:

converter_flags = {
    "optimizations": [tf.lite.Optimize.DEFAULT],
    "target_spec": {
        "supported_types": [tf.float16]
    }
}

edge_model = ai_edge_torch.convert(
    model, sample_inputs, _ai_edge_converter_flags=converter_flags
)

This should correctly trigger the float16 quantization.

[Bludator]

With those flags it has started doing something but it crashes nevertheless with these logs:

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1766583376.590961   11490 tf_tfl_flatbuffer_helpers.cc:364] Ignored output_format.
W0000 00:00:1766583376.591000   11490 tf_tfl_flatbuffer_helpers.cc:367] Ignored drop_control_dependency.
I0000 00:00:1766583376.591660   11490 reader.cc:83] Reading SavedModel from: /tmp/tmp2_byoph2
I0000 00:00:1766583376.592002   11490 reader.cc:52] Reading meta graph with tags { serve }
I0000 00:00:1766583376.592015   11490 reader.cc:147] Reading SavedModel debug info (if present) from: /tmp/tmp2_byoph2
I0000 00:00:1766583376.594024   11490 mlir_graph_optimization_pass.cc:437] MLIR V1 optimization pass is not enabled
I0000 00:00:1766583376.594379   11490 loader.cc:236] Restoring SavedModel bundle.
I0000 00:00:1766583376.608078   11490 loader.cc:220] Running initialization op on SavedModel bundle at path: /tmp/tmp2_byoph2
I0000 00:00:1766583376.613206   11490 loader.cc:471] SavedModel load for tags { serve }; Status: success: OK. Took 21562 microseconds.
Illegal instruction (core dumped)