wehub-resource-sync
cddb07a176
build container image / cpu (push) Waiting to run
build container image / cuda (push) Waiting to run
build container image / rocm (push) Waiting to run
frontend tests / frontend-tests (push) Waiting to run
openapi checks / openapi-checks (push) Waiting to run
python tests / py3.12: macos-default (push) Waiting to run
python tests / py3.11: windows-cpu (push) Waiting to run
python tests / py3.12: windows-cpu (push) Waiting to run
python tests / py3.11: linux-cpu (push) Waiting to run
python tests / py3.12: linux-cpu (push) Waiting to run
typegen checks / typegen-checks (push) Waiting to run
uv lock checks / uv-lock-checks (push) Waiting to run
frontend checks / frontend-checks (push) Waiting to run
lfs checks / lfs-check (push) Waiting to run
python checks / python-checks (push) Waiting to run
python tests / py3.11: macos-default (push) Waiting to run
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
74 lines
3.2 KiB
Python
74 lines
3.2 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from invokeai.backend.model_manager.load.optimizations import _no_op, skip_torch_weight_init
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
["torch_module", "layer_args"],
|
|
[
|
|
(torch.nn.Linear, {"in_features": 10, "out_features": 20}),
|
|
(torch.nn.Conv1d, {"in_channels": 10, "out_channels": 20, "kernel_size": 3}),
|
|
(torch.nn.Conv2d, {"in_channels": 10, "out_channels": 20, "kernel_size": 3}),
|
|
(torch.nn.Conv3d, {"in_channels": 10, "out_channels": 20, "kernel_size": 3}),
|
|
(torch.nn.Embedding, {"num_embeddings": 10, "embedding_dim": 10}),
|
|
],
|
|
)
|
|
def test_skip_torch_weight_init_linear(torch_module, layer_args):
|
|
"""Test the interactions between `skip_torch_weight_init()` and various torch modules."""
|
|
seed = 123
|
|
|
|
# Initialize a torch layer *before* applying `skip_torch_weight_init()`.
|
|
reset_params_fn_before = torch_module.reset_parameters
|
|
torch.manual_seed(seed)
|
|
layer_before = torch_module(**layer_args)
|
|
|
|
# Initialize a torch layer while `skip_torch_weight_init()` is applied.
|
|
with skip_torch_weight_init():
|
|
reset_params_fn_during = torch_module.reset_parameters
|
|
torch.manual_seed(123)
|
|
layer_during = torch_module(**layer_args)
|
|
|
|
# Initialize a torch layer *after* applying `skip_torch_weight_init()`.
|
|
reset_params_fn_after = torch_module.reset_parameters
|
|
torch.manual_seed(123)
|
|
layer_after = torch_module(**layer_args)
|
|
|
|
# Check that reset_parameters is skipped while `skip_torch_weight_init()` is active.
|
|
assert reset_params_fn_during == _no_op
|
|
assert not torch.allclose(layer_before.weight, layer_during.weight)
|
|
if hasattr(layer_before, "bias"):
|
|
assert not torch.allclose(layer_before.bias, layer_during.bias)
|
|
|
|
# Check that the original behavior is restored after `skip_torch_weight_init()` ends.
|
|
assert reset_params_fn_before is reset_params_fn_after
|
|
assert torch.allclose(layer_before.weight, layer_after.weight)
|
|
if hasattr(layer_before, "bias"):
|
|
assert torch.allclose(layer_before.bias, layer_after.bias)
|
|
|
|
|
|
def test_skip_torch_weight_init_restores_base_class_behavior():
|
|
"""Test that `skip_torch_weight_init()` correctly restores the original behavior of torch.nn.Conv*d modules. This
|
|
test was created to catch a previous bug where `reset_parameters` was being copied from the base `_ConvNd` class to
|
|
its child classes (like `Conv1d`).
|
|
"""
|
|
with skip_torch_weight_init():
|
|
# There is no need to do anything while the context manager is applied, we're just testing that the original
|
|
# behavior is restored correctly.
|
|
pass
|
|
|
|
# Mock the behavior of another library that monkey patches `torch.nn.modules.conv._ConvNd.reset_parameters` and
|
|
# expects it to affect all of the sub-classes (e.g. `torch.nn.Conv1D`, `torch.nn.Conv2D`, etc.).
|
|
called_monkey_patched_fn = False
|
|
|
|
def monkey_patched_fn(*args, **kwargs):
|
|
nonlocal called_monkey_patched_fn
|
|
called_monkey_patched_fn = True
|
|
|
|
saved_fn = torch.nn.modules.conv._ConvNd.reset_parameters
|
|
torch.nn.modules.conv._ConvNd.reset_parameters = monkey_patched_fn
|
|
_ = torch.nn.Conv1d(10, 20, 3)
|
|
torch.nn.modules.conv._ConvNd.reset_parameters = saved_fn
|
|
|
|
assert called_monkey_patched_fn
|