feat: implementing Kernel Inception Metric #65 : #81

dhruvmalik007 · 2025-05-06T09:14:47Z

Author:

Intro

This pull request introduces a new implementation of the Kernel Inception Distance(KID) metric for assessing the quality of generated images and includes comprehensive unit tests to validate its functionality. The most important changes are the addition of the KernelInceptionMetric class, which computes the KID using a polynomial kernel, and the creation of a robust test suite to ensure the metric's correctness, stability, and sensitivity to various scenarios.

Changes :

image_metrics.py: Added the KernelInceptionDistanceMetric class, which computes the Kernel Inception Distance (KID) using a polynomial kernel. This includes methods for feature accumulation, MMD computation, and KID calculation with configurable parameters like subsets, subset size, kernel degree, gamma, and coefficient.
nnx_metrics.py: defining the wrapper method for the support to the NNX library to import the method directly
metrax_test.py: defined the random test case.
__init__.py: and adding the import alias for the KID metric.

Unit Tests:

src/metrax/image_metrics_test.py: Added a comprehensive test suite for KernelInceptionMetric with the following test cases:
- Validation of default parameter functionality.
- Handling of invalid parameters and edge cases.
- Behavior with small sample sizes and identical feature sets.
- Sensitivity to mode collapse, outliers, and different subset configurations.
- Differentiation between similar and dissimilar feature distributions.

- writing initial version of the kernel incpetion disrance metric - writing the test cases to cover edge case for classification.

google-cla · 2025-05-06T09:14:51Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

jshin1394 · 2025-05-06T21:16:25Z

Hi Dhruv, thank you so much for your contribution to the Metrax repo! The first iteration of your PR looks very promising! :)

A few things:

Could you please sign the Google Contributor License Agreement(CLA) as suggested by the google-cla bot?
Here is an example PR which adds a new metric to Metrax(https://github.com/google/metrax/pull/79/files). It also adds the metric to metrax_test in order to test its jittability and adds the metrax to metrax/nnx in order to support NNX users.
clu_metric.Metric interface has a list of functions that must be implemented and merge is one of them(https://github.com/google/metrax/pull/79/files).

Thank you so much! :)

- adding the import modules and name change - NNX class wrapper description - implementing merge and empty method as required by the clu_metrics.Metrics method - and also updating the tests in the metrax_test along with the unittests.

jshin1394

Thank you so much Dhruv!

src/metrax/image_metrics.py

src/metrax/image_metrics_test.py

jshin1394 · 2025-05-07T20:50:30Z

src/metrax/image_metrics.py

+
+
+@flax.struct.dataclass
+class KernelInceptionDistanceMetric(clu_metrics.Metric):


what are your thoughts on having this inherit base.Average from the metrax package? It will keep track of the sum of total KID and the sum of total count and compute will yield the average KID of the data.

indeed I have done that and during my prompting , I realised to implement it in the following function:

@classmethod def from_model_output( cls, real_features: jax.Array, fake_features: jax.Array, subsets: int = KID_DEFAULT_SUBSETS, subset_size: int = KID_DEFAULT_SUBSET_SIZE, degree: int = KID_DEFAULT_DEGREE, gamma: float = KID_DEFAULT_GAMMA, coef: float = KID_DEFAULT_COEF, ): # checks for the valid inputs if subsets <= 0 or subset_size <= 0 or degree <= 0 or (gamma is not None and gamma <= 0) or coef <= 0: raise ValueError("All parameters must be positive and non-zero.") # Compute KID for this batch if real_features.shape[0] < subset_size or fake_features.shape[0] < subset_size: raise ValueError("Subset size must be smaller than the number of samples.") master_key = random.PRNGKey(42) kid_scores = [] for i in range(subsets): key_real, key_fake = random.split(random.fold_in(master_key, i)) real_indices = random.choice(key_real, real_features.shape[0], (subset_size,), replace=False) fake_indices = random.choice(key_fake, fake_features.shape[0], (subset_size,), replace=False) f_real_subset = real_features[real_indices] f_fake_subset = fake_features[fake_indices] kid = cls.__compute_mmd_static(f_real_subset, f_fake_subset, degree, gamma, coef) kid_scores.append(kid) kid_mean = jnp.mean(jnp.array(kid_scores)) # Accumulate sum and count for averaging return cls( total=kid_mean, count=1.0, subsets=subsets, subset_size=subset_size, degree=degree, gamma=gamma, coef=coef, )

- `image_metrics.py`/ `metrax_test.py`: changing name along with defining the private functions and other details. - test cases added to compare with the torchmetrics implementation. - installing the requisite packages to support torchmetrics comparison

dhruvmalik007 · 2025-05-08T19:03:40Z

Thank you so much Dhruv!

your welcome and appreciated your help for resolving issues.

jshin1394

Thank you :)

jshin1394 · 2025-05-08T19:53:36Z

src/metrax/__init__.py

 from metrax import ranking_metrics
 from metrax import regression_metrics
-
+from metrax import image_metrics


let's remove this line since image_metrics was added as part of previous PR in line 17.

src/metrax/image_metrics.py

jshin1394 · 2025-05-08T19:57:47Z

src/metrax/image_metrics.py

+        return result
+
+
+    def merge(self, other: "KernelInceptionDistance") -> "KernelInceptionDistance":


now that we are inheriting base.Average, we shouldn't need to implement compute and merge.

jshin1394 · 2025-05-08T19:57:59Z

src/metrax/image_metrics.py

+        )
+
+    @classmethod
+    def empty(cls) -> "KernelInceptionDistance":


now that we are inheriting base.Average, we shouldn't need to implement empty

src/metrax/image_metrics.py

…ocal: - again renamed to KID - restructuring the test util function to image_metrics.py .

…uvmalik007/metrax into local branch

@jshin1394

- minor function name changes . - removing the redundant methods implemented for clu_metrics - also merging the @jshin1394 changes and then merging our works.

jshin1394

Thank you so much :)

jshin1394 · 2025-05-09T17:46:56Z

src/metrax/__init__.py

 SSIM = image_metrics.SSIM
 WER = nlp_metrics.WER
-
+KID = image_metrics.KID


nit: let's keep alphabetical order.

src/metrax/__init__.py

src/metrax/image_metrics.py

src/metrax/image_metrics_test.py

jshin1394

Thank you so much :)

jshin1394 · 2025-05-15T15:24:44Z

src/metrax/image_metrics.py

+        cls,
+        real_features: jax.Array,
+        fake_features: jax.Array,
+        subsets: int = KID_DEFAULT_SUBSETS,


can we not define member variables and just replace with default values here?
subsets: int = 1000,
subset_size: int = 100,
degree: int = 3,
gamma: float = 1.0,
coef: float = 1.0,

Sounds good. It was just in case if some model_eval.py evaluation script might import this parameter .

src/metrax/image_metrics.py

src/metrax/nnx/__init__.py

src/metrax/nnx/nnx_metrics.py

src/metrax/image_metrics.py

- Resolve KID metric position in nnx - Removing the constants - Adding doc-string for model_output with description.

dhruvmalik007 · 2025-06-02T16:54:56Z

hi @jshin1394 , apologies of not notifying on this branch, let me know if are there any remaining issues to be resolved .

jshin1394 · 2025-06-02T19:33:27Z

hi @jshin1394 , apologies of not notifying on this branch, let me know if are there any remaining issues to be resolved .

Hi Dhruv, are we assuming that users will call metrax.KID with extracted features only? It seems like torchmetrics.KID takes in as input the real image and the generate image but metrax.KID takes in real_features and generated_features.

…d image.

dhruvmalik007 · 2025-06-04T04:12:12Z

hi @jshin1394 , apologies of not notifying on this branch, let me know if are there any remaining issues to be resolved .

Hi Dhruv, are we assuming that users will call metrax.KID with extracted features only? It seems like torchmetrics.KID takes in as input the real image and the generate image but metrax.KID takes in real_features and generated_features.

Yes you're right . initially I have taken the I{Real} and I{Fake} as in the reference implementation here does indeed take the generated features . I have done the changes accordingly .

jshin1394 · 2025-06-10T22:56:38Z

hi @jshin1394 , apologies of not notifying on this branch, let me know if are there any remaining issues to be resolved .

Hi Dhruv, are we assuming that users will call metrax.KID with extracted features only? It seems like torchmetrics.KID takes in as input the real image and the generate image but metrax.KID takes in real_features and generated_features.

Yes you're right . initially I have taken the I{Real} and I{Fake} as in the reference implementation here does indeed take the generated features . I have done the changes accordingly .

I think the original implementation takes in images as input and calls this model(https://github.com/Lightning-AI/torchmetrics/blob/74bdb26b35515c36334f59d654f0f4d8a12d5b59/src/torchmetrics/image/kid.py#L211) to get the features. If the implementation does indeed require pretrained model, I will go ahead and come up with a design for this.

feat: implementing google#65 :

2af982a

- writing initial version of the kernel incpetion disrance metric - writing the test cases to cover edge case for classification.

dhruvmalik007 mentioned this pull request May 6, 2025

Add Kernel Inception Distance (KID) metric #65

Open

dhruvmalik007 changed the title ~~feat: implementing #65 :~~ feat: implementing Kernel Inception Metric #65 : May 7, 2025

dhruvmalik007 added 2 commits May 7, 2025 09:45

oops: forgot the contribution of the merge and empty methods

c4bea1b

jshin1394 reviewed May 7, 2025

View reviewed changes

dhruvmalik007 added 2 commits May 8, 2025 17:44

removing the logging (added to test the errors before)

87b736a

Merge branch 'main' into dhruvmalik007/add-kid-metric

f13d26e

jshin1394 reviewed May 8, 2025

View reviewed changes

dhruvmalik007 added 5 commits May 8, 2025 21:28

stashing local changes before pulling the origin branch changesinto l…

c1c30be

…ocal: - again renamed to KID - restructuring the test util function to image_metrics.py .

Merge branch 'dhruvmalik007/add-kid-metric' of https://github.com/dhr…

24b16d0

…uvmalik007/metrax into local branch

refactor:

88f9ec6

- minor function name changes . - removing the redundant methods implemented for clu_metrics - also merging the @jshin1394 changes and then merging our works.

all test passing

6bdd189

refactor: test uncluding ersthwhile empty and merge functions removed

4cb3d65

jshin1394 reviewed May 9, 2025

View reviewed changes

dhruvmalik007 added 2 commits May 10, 2025 00:34

Merge branch 'main' into dhruvmalik007/add-kid-metric

f9d5954

add the change suggestions @jshin1394 .

965eab5

jshin1394 reviewed May 15, 2025

View reviewed changes

dhruvmalik007 added 4 commits May 15, 2025 16:15

minor refactor:

075e68d

- Resolve KID metric position in nnx - Removing the constants - Adding doc-string for model_output with description.

merging from the remote to the local

878876c

Merge branch 'main' into dhruvmalik007/add-kid-metric

d40ee0c

Merge branch 'main' into dhruvmalik007/add-kid-metric

d3df741

feat: replacing the KID metrics computation on actual and fake derive…

8969ce9

…d image.

Merge branch 'main' into dhruvmalik007/add-kid-metric

74cf343



		@flax.struct.dataclass
		class KernelInceptionDistanceMetric(clu_metrics.Metric):

		return result


		def merge(self, other: "KernelInceptionDistance") -> "KernelInceptionDistance":

feat: implementing Kernel Inception Metric #65 : #81

Are you sure you want to change the base?

feat: implementing Kernel Inception Metric #65 : #81

Uh oh!

Conversation

dhruvmalik007 commented May 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Author:

Intro

Changes :

Unit Tests:

Uh oh!

google-cla bot commented May 6, 2025

Uh oh!

jshin1394 commented May 6, 2025

Uh oh!

jshin1394 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dhruvmalik007 commented May 8, 2025

Uh oh!

jshin1394 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jshin1394 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jshin1394 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dhruvmalik007 commented Jun 2, 2025

Uh oh!

jshin1394 commented Jun 2, 2025

Uh oh!

dhruvmalik007 commented Jun 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jshin1394 commented Jun 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

dhruvmalik007 commented May 6, 2025 •

edited

Loading

dhruvmalik007 commented Jun 4, 2025 •

edited

Loading