apple / apple/coremltools

Tensorflow BatchNormalization conversion error: mean type are not broadcast compatible

Open
#2,462 1 comment 1 reaction 0 assignees View on GitHub
bug
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🐞Describing the bug
I'm converting a tensorflow keras model containing many BatchNormalization layers, and i kept getting a strange error about broadcasting errors, with a nonsensical shape for the mean variable, which kept changing every time i ran the code.
To hunt down the bug i created a keras model that was a single BatchNormalization layer, taking my input batch shape, which failed to convert to coreml with the same error as before.
The strange thing is that the shape it is trying to broadcast keeps changing, without changing any code.
This smells to me like a faulty memory read somewhere. Perhaps it is a wrong index read, which is supported by the fact that if i change the default axis in the BatchNormalization layer from -1 to 3, then everything works.
However i am then normalizing on the wrong axis, as my input has 5 dimension, and thus axis -1 should be 4 not 3.
If i try with axis=4 then it fails again with the same error.
Weirdly it also fails on axis 0 and 1, but not axis 2 and 3.

I have also checked if it is the tensor shape that is the culprit, by transposing my last axis to index 2, and there it has no issues.

## Stack Trace
- If applicable, please paste the complete stack trace.
```
Running TensorFlow Graph Passes: 100%|██████████| 6/6 [00:00<00:00, 92.50 passes/s]
Converting TF Frontend ==> MIL Ops: 100%|██████████| 7/7 [00:00<00:00, 4478.36 ops/s]
Running MIL frontend_tensorflow2 pipeline: 100%|██████████| 7/7 [00:00<00:00, 7909.52 passes/s]
Running MIL default pipeline: 100%|██████████| 87/87 [00:00<00:00, 5561.30 passes/s]
Running MIL backend_mlprogram pipeline: 100%|██████████| 12/12 [00:00<00:00, 11767.98 passes/s]
loc("tensor test_fail_batch_normalization_FusedBatchNormV3 = batch_norm(beta = tensor([0, 0]), epsilon = fp32(0.00100000005), gamma = tensor([1, 1]), mean = tensor([0, 0]), variance = tensor([1, 1]), x = transpose_0)[milId = uint64(1), name = string(\22test_fail_batch_normalization_FusedBatchNormV3\22)]; - /private/var/folders/m6/pjcdf66928536cc45np45sqw0000gn/T/tmpr0zgd3rp.mlmodelc/model.mil":12:12): error: output type 'tensor<8x2x128x64x32xf32>' and mean type 'tensor<1x0x1x1x1975375664xf32>' are not broadcast compatible
LLVM ERROR: Failed to infer result type(s).
```
Next run the broadcast shape is different without changing the input:
```
error: output type 'tensor<8x2x128x64x32xf32>' and mean type 'tensor<1x0x1x1x2019839792xf32>' are not broadcast compatible
```
For comparison
First run:```'tensor<1x0x1x1x1975375664xf32>'```
Second run: ```'tensor<1x0x1x1x2019839792xf32>'```

## To Reproduce
```
import tensorflow as tf
import coremltools as ct
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.layers import Input

def make_failing_dummy_model(input_batch_shape):
inputs = Input(batch_shape=input_batch_shape)
a = BatchNormalization(axis=-1)
x = a(inputs)

model = tf.keras.Model(inputs=inputs, outputs=x, name="test_fail")
return model

def convert_model(model, input_batch_shape):
input_shape = ct.TensorType(name="input_1", shape=input_batch_shape,
dtype=ct.converters.mil.input_types.types.fp32)

coreml_model = ct.convert(model,
inputs=[input_shape],
source='tensorflow',
convert_to="mlprogram",
compute_precision=ct.precision.FLOAT32,
minimum_deployment_target=ct.target.macOS12)

input_batch_shape = (8,128,64,32,2)
model = make_failing_dummy_model(input_batch_shape)
convert_model(model, input_batch_shape)
```

## System environment (please complete the following information):
- coremltools version: 8.2
- OS (e.g. MacOS version or Linux type): MacOS 15.3.1
- Any other relevant version information (e.g. PyTorch or TensorFlow version):
- tensorflow-macos: 2.12.0
- tensorflow-metal: 1.2.0
- keras: 2.12.0

## Additional context
Here is my entire testing code, to show the three scenarios i described above:
```
# Paste Python code snippet here, complete with any required import statements.
import traceback
import numpy as np
import tensorflow as tf
import coremltools as ct
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.layers import Input

def make_failing_dummy_model(input_batch_shape):
inputs = Input(batch_shape=input_batch_shape)
a = BatchNormalization(axis=-1)
x = a(inputs)

model = tf.keras.Model(inputs=inputs, outputs=x, name="test_fail")
return model

def make_succeeding_wrong_dummy_model(input_batch_shape):
inputs = Input(batch_shape=input_batch_shape)
a = BatchNormalization(axis=3)
x = a(inputs)
model = tf.keras.Model(inputs=inputs, outputs=x, name="test_success_wrong")
return model

def make_succeeding_transposed_dummy_model(input_batch_shape):
inputs = Input(batch_shape=input_batch_shape)
t = tf.transpose(inputs, (0,1,4,2,3))
a = BatchNormalization(axis=2)
x = a(t)

model = tf.keras.Model(inputs=inputs, outputs=x, name="test_success_transpose")
return model

def convert_model(model, input_batch_shape):
input_shape = ct.TensorType(name="input_1", shape=input_batch_shape,
dtype=ct.converters.mil.input_types.types.fp32)

coreml_model = ct.convert(model,
inputs=[input_shape],
source='tensorflow',
convert_to="mlprogram",
compute_precision=ct.precision.FLOAT32,
minimum_deployment_target=ct.target.macOS12)
def test_failing(input_batch_shape):
#This try except does nothing, as the program crashes hard
try:
model = make_failing_dummy_model(input_batch_shape)
convert_model(model, input_batch_shape)
except:
traceback.print_tb()
traceback.print_exception()
pass

def test_succeed_transpose(input_batch_shape):
model = make_succeeding_transposed_dummy_model(input_batch_shape)
convert_model(model, input_batch_shape)
del model

def test_succeed_wrong(input_batch_shape):
model = make_succeeding_wrong_dummy_model(input_batch_shape)
convert_model(model, input_batch_shape)
del model

def test_three_methods():
input_batch_shape = (8, 128, 64, 32, 2)

test_succeed_wrong(input_batch_shape)
print("Succeeded in conversion with axis=3")
tf.keras.backend.clear_session()
test_succeed_transpose(input_batch_shape)
print("Succeeded in conversion with transpose and axis=2")
tf.keras.backend.clear_session()
test_failing(input_batch_shape)
print("Test failed (This is never printed as the program hard crashes before)")

test_three_methods()
```

Contributor guide

Open the contributing guide

Research direction

Run the provided minimal reproduction with coremltools 8.2 and the listed TensorFlow versions, then inspect the TensorFlow BatchNormalization conversion for FusedBatchNormV3. Compare axis=-1 with the working axis and transpose cases, and repeat conversions to confirm whether the generated mean shape changes. Done means the five-dimensional axis=-1 model converts without a broadcast error or hard crash.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, tensorflow
Domain
machine-learning, tooling
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.