refactor unet_1d tests#13898
Conversation
|
@askserge could you please review the PR? (testing) |
![sergereview[bot]](https://avatars.githubusercontent.com/in/3673521?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤗 Serge says:
The migration to the mixin-based test structure is mostly correct, but there are issues with the output_shape property values.
Correctness
-
UNetRLModelTesterConfig.output_shapeis set to(4, 14, 1)but the RL value-function model actually outputs(batch_size, 1)— i.e.(4, 1)for batch_size=4. The old test explicitly checkedtorch.Size((inputs_dict["sample"].shape[0], 1)). This is wrong regardless of whether the convention is per-sample or full-batch shape. It doesn't cause failures today because the basetest_outputhas an operator-precedence bug (== … or self.output_shapeis always truthy), and the customtest_outputoverride doesn't useself.output_shape. But it will break when the base class is fixed, and it's misleading to readers. -
UNet1DModelTesterConfig.output_shapeis(4, 14, 16)(includes batch dim), while the repo convention in other migrated tests (e.g.AutoencoderKLTesterConfig) is to use per-sample shape without the batch dimension —(3, 32, 32)not(4, 3, 32, 32). Should be(14, 16)to match the convention.
Minor
- In
TestUNetRLModel.test_output(line 173),self.get_dummy_inputs()is called a second time just to readshape[0](which is the hardcodedbatch_size = 4). This allocates a full set of tensors unnecessarily.
model: claude-opus-4-6 · 28 LLM turns · 30 tool calls · 143.1s · 627103 in / 6681 out tokens
| @property | ||
| def output_shape(self): | ||
| def output_shape(self) -> tuple: | ||
| return (4, 14, 16) |
There was a problem hiding this comment.
output_shape includes the batch dimension (4, 14, 16), but the repo convention (see e.g. AutoencoderKLTesterConfig) is to use the per-sample shape without the batch dim. Should be (14, 16) to stay consistent.
| return (4, 14, 16) | |
| return (14, 16) |
| @property | ||
| def output_shape(self): | ||
| def output_shape(self) -> tuple: | ||
| return (4, 14, 1) |
There was a problem hiding this comment.
output_shape is (4, 14, 1) but the RL value-function model outputs shape (batch_size, 1). The old test explicitly asserted torch.Size((batch_size, 1)). This value is wrong regardless of convention (per-sample it should be (1,), full-batch it should be (4, 1)).
The base test_output currently has an operator-precedence bug that makes the assertion vacuously true, and the custom test_output override doesn't reference self.output_shape, so this doesn't fail today — but it will break when the base class is fixed.
| return (4, 14, 1) | |
| return (1,) |
| with torch.no_grad(): | ||
| output = model(**self.get_dummy_inputs()).sample | ||
|
|
||
| assert output.shape == (self.get_dummy_inputs()["sample"].shape[0], 1), "Input and output shapes do not match" |
There was a problem hiding this comment.
Nit: self.get_dummy_inputs() is called a second time just to read shape[0], which is always the hardcoded batch_size = 4. This unnecessarily allocates a full set of tensors. Consider storing the inputs in a local variable:
| assert output.shape == (self.get_dummy_inputs()["sample"].shape[0], 1), "Input and output shapes do not match" | |
| inputs = self.get_dummy_inputs() | |
| assert output.shape == (inputs["sample"].shape[0], 1), "Input and output shapes do not match" |
(and pass inputs to the model call above as well)
* refactor unet_1d tests * use per-sample output_shape for unet_1d tests --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2 participants
What does this PR do?
Part of the ongoing modeling-test migration (following #13369 and #13153). Migrates the UNet1DModel test suites (the standard UNet and the RL value-function variant) to the mixin-based structure (Config + ModelTesterMixin).
UNet1DModel has no attention processors and doesn't support gradient checkpointing, so only ModelTesterMixin applies. The hub/pretrained integration tests are kept.
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
@sayakpaul @DN6