Как заставить работать пользовательскую функцию потерь с большим количеством переменных входов ⇐ Python

[Anonymous]

Я создаю или пытаюсь создать модель, которая сможет оценить внутренние параметры камеры. Для этого я хочу использовать пользовательскую функцию потерь, которая использует оцененные внутренние параметры камеры, а затем неискажает изображение и сравнивает его с изображением глубины (из того же места, полученным из облака точек лидара). Но функция потерь выдает следующую ошибку:

Код: Выделить всё

Node: 'gradient_tape/model/pose_loss_layer/undistort_layer/map/while/gradients/model/pose_loss_layer/undistort_layer/map/while/strided_slice_2_grad/StridedSliceGrad'
shape of dy was [] instead of [5]
[[{{node gradient_tape/model/pose_loss_layer/undistort_layer/map/while/gradients/model/pose_loss_layer/undistort_layer/map/while/strided_slice_2_grad/StridedSliceGrad}}]] [Op:__inference_train_function_2727]

Вот код, который я использую для проверки функции потерь (это только для проверки функции потерь, в другой модели я использую больше слоев и все такое, но эта модель дает то же самое ошибка).

Код: Выделить всё

class UndistortLayer(Layer):
def __init__(self):
super(UndistortLayer, self).__init__()

def call(self, inputs):
camera_images, intrinsic_params, distortion_coeffs = inputs

def undistort_image(camera_image, intrinsic_params, distortion_coeffs):
# Create intrinsic matrix
intrinsic_matrix = create_intrinsic_matrix(intrinsic_params)

# Undistort image using OpenCV
undistorted_image = cv2.undistort(
camera_image.numpy().astype(np.uint8),
intrinsic_matrix,
np.array(distortion_coeffs, dtype=np.float32)
)
return undistorted_image.astype(np.float32)

# Apply the undistortion function using tf.py_function for each image
undistorted_images = tf.map_fn(
lambda i: tf.py_function(
undistort_image,
[camera_images[i], intrinsic_params[i], distortion_coeffs[i]],
Tout=tf.float32
),
tf.range(tf.shape(camera_images)[0]),  # Loop over the batch
dtype=tf.float32
)
undistorted_images.set_shape([None, 480, 640, 3])
return undistorted_images

class PoseLossLayer(tf.keras.layers.Layer):
def call(self, inputs):
output_intrinsic, camera_input, depth_input = inputs

camera_images = camera_input  # Shape: (batch_size, 480, 640, 3)
depth_images = depth_input    # Shape: (batch_size, 480, 640)

# Retrieve intrinsic parameters and distortion coefficients from predictions
intrinsic_params = output_intrinsic[..., :3]  # fx, cx, cy
distortion_coeffs = output_intrinsic[..., 3:]  # distortion coefficients

# Ensure UndistortLayer supports backpropagation
undistorted_camera_images = UndistortLayer()([camera_images, intrinsic_params, distortion_coeffs])

# Convert to grayscale correctly
undistorted_camera_images_gray = tf.image.rgb_to_grayscale(undistorted_camera_images)

# Calculate SSIM between undistorted camera images and depth images
ssim_loss = tf.image.ssim(undistorted_camera_images_gray, depth_images, max_val=1.0)

# Compute the mean SSIM loss to ensure a scalar is returned
mean_ssim_loss = tf.reduce_mean(ssim_loss)

# Return a scalar loss value for backpropagation
return 1 - mean_ssim_loss

def create_model():
# Define flexible input layers
camera_input = Input(shape=(None, None, 3), name='camera_image')
depth_input = Input(shape=(None, None, 1), name='depth_image')

# Process depth image to have 3 channels
depth_input_3ch = Lambda(lambda x:  tf.image.grayscale_to_rgb(x))(depth_input)

# Concatenate camera and depth images
concatenated_features = Concatenate(name='RGB_Depth_concatenate')([camera_input, depth_input_3ch])

# Add Global Average Pooling to reduce spatial dimensions
pooled_features = GlobalAveragePooling2D()(concatenated_features)

# Dense layers
x = Dense(64, activation='relu')(pooled_features)
x = Dense(32, activation='relu')(x)
output_intrinsic = Dense(8, activation='linear', name='intrinsic_parameters')(x)

# PoseLossLayer to compute loss inside the graph
pose_loss = PoseLossLayer()([output_intrinsic, camera_input, depth_input])

# Define the model
model = Model(inputs=[camera_input, depth_input], outputs=output_intrinsic)

# Add loss to model
model.add_loss(pose_loss)

# Compile the model (without loss, as it's added via add_loss)
model.compile(optimizer='adam')

return model

model.fit(
x=[camera_images_array, depth_images_array],
y=Parameters_Short,
epochs=1,
batch_size=5,
callbacks=[CustomCallback()]
)

Надеюсь, кто-нибудь знает проблему или лучший способ добавить функцию потерь.
пробовал изменить формат возвращаемого значения, пытался добавить его в model.compile, но ничего не работает

Подробнее здесь: https://stackoverflow.com/questions/790 ... nputs-work