-
Notifications
You must be signed in to change notification settings - Fork 4.9k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Test ResnetBlock2D #1850
Test ResnetBlock2D #1850
Conversation
The documentation is not available anymore as the PR was closed or merged. |
tests/test_layers_utils.py
Outdated
def test_resnet_default(self): | ||
torch.manual_seed(0) | ||
sample = torch.randn(1, 32, 64, 64) | ||
temb = torch.randn(1, 128) | ||
resnet_block = ResnetBlock2D(in_channels=32, temb_channels=128) | ||
with torch.no_grad(): | ||
output_tensor = resnet_block(sample, temb) | ||
|
||
assert output_tensor.shape == (1, 32, 64, 64) | ||
output_slice = output_tensor[0, -1, -3:, -3:] | ||
expected_slice = torch.tensor([-1.9010, -0.2974, -0.8245, -1.3533, 0.8742, -0.9645, -2.0584, 1.3387, -0.4746]) | ||
assert torch.allclose(output_slice.flatten(), expected_slice, atol=1e-3) | ||
|
||
def test_restnet_with_use_in_shortcut(self): | ||
torch.manual_seed(0) | ||
sample = torch.randn(1, 32, 64, 64) | ||
temb = torch.randn(1, 128) | ||
resnet_block = ResnetBlock2D(in_channels=32, temb_channels=128, use_in_shortcut=True) | ||
with torch.no_grad(): | ||
output_tensor = resnet_block(sample, temb) | ||
|
||
assert output_tensor.shape == (1, 32, 64, 64) | ||
output_slice = output_tensor[0, -1, -3:, -3:] | ||
expected_slice = torch.tensor([0.2226, -1.0791, -0.1629, 0.3659, -0.2889, -1.2376, 0.0582, 0.9206, 0.0044]) | ||
assert torch.allclose(output_slice.flatten(), expected_slice, atol=1e-3) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks good! We should also test this for GPU, this can be done using the torch_device
. Check the AttentionBlockTests
for reference. Let me know if you need any help.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@patil-suraj Added and tested on both GPU and MPS. Lmk what else is needed. Thanks!
There are a few fast tests failing though -- those are due to test_pipelines.py::CustomPipelineTests
, which is irrelevant to the changes here and has been failing from the main branch.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@patil-suraj let me know if we can merge this. Thanks ~
This looks super cool - thanks a lot @hchings :-) |
What does this PR do?
Add test cases for Resnet block. Issue: #200