diff --git a/CHANGELOG.md b/CHANGELOG.md index 6217e76d..14286b4f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,10 @@ changelog does not include internal changes that do not affect the user. ### Fixed +- Fixed `.to()` not moving the tensors given to the `Constant`, `STCH`, `PBI` and `COSMOS` + scalarizers (`weights` and `reference`). These tensors are now registered as non-persistent + buffers, so they follow the dtype and device of the scalarizer without being added to its + `state_dict`. `FAMO`'s `min_losses` buffer is also no longer in its `state_dict`. - Fixed `.to()` not moving the tensors given to `Constant`, `ConstantWeighting`, `UPGrad`, `UPGradWeighting`, `DualProj`, `DualProjWeighting`, `AlignedMTL`, `AlignedMTLWeighting`, `ConFIG`, `SDMGradWeighting`, `CRMOGMWeighting` and `GradDrop` (`weights`, `pref_vector`, diff --git a/src/torchjd/scalarization/_constant.py b/src/torchjd/scalarization/_constant.py index 5223f289..0d46781f 100644 --- a/src/torchjd/scalarization/_constant.py +++ b/src/torchjd/scalarization/_constant.py @@ -14,9 +14,11 @@ class Constant(Scalarizer): passed at call time. """ + weights: Tensor + def __init__(self, weights: Tensor) -> None: super().__init__() - self.weights = weights + self.register_buffer("weights", weights, persistent=False) def forward(self, values: Tensor, /) -> Tensor: if values.shape != self.weights.shape: diff --git a/src/torchjd/scalarization/_cosmos.py b/src/torchjd/scalarization/_cosmos.py index 2de5077c..9135ed7c 100644 --- a/src/torchjd/scalarization/_cosmos.py +++ b/src/torchjd/scalarization/_cosmos.py @@ -42,6 +42,8 @@ class COSMOS(Scalarizer): objective. """ + weights: Tensor + def __init__(self, lambda_: float, weights: Tensor) -> None: if lambda_ < 0.0: raise ValueError( @@ -50,7 +52,7 @@ def __init__(self, lambda_: float, weights: Tensor) -> None: super().__init__() self.lambda_ = lambda_ - self.weights = weights + self.register_buffer("weights", weights, persistent=False) def forward(self, values: Tensor, /) -> Tensor: check_same_shape_as_values(self.weights, values, "weights") diff --git a/src/torchjd/scalarization/_famo.py b/src/torchjd/scalarization/_famo.py index 7fca1b41..9d4eb294 100644 --- a/src/torchjd/scalarization/_famo.py +++ b/src/torchjd/scalarization/_famo.py @@ -128,7 +128,7 @@ def __init__( f"Parameter `min_losses` should have shape {tuple(self._w.shape)} (matching the " f"shape of the logits). Found `min_losses.shape = {tuple(min_losses.shape)}`." ) - self.register_buffer("min_losses", min_losses) + self.register_buffer("min_losses", min_losses, persistent=False) self.lr = lr self.weight_decay = weight_decay diff --git a/src/torchjd/scalarization/_pbi.py b/src/torchjd/scalarization/_pbi.py index fe8f5025..fb6734a9 100644 --- a/src/torchjd/scalarization/_pbi.py +++ b/src/torchjd/scalarization/_pbi.py @@ -51,14 +51,17 @@ class PBI(Scalarizer): :math:`10^{-6}` at that point and is negligible elsewhere. """ + weights: Tensor + reference: Tensor | None + def __init__(self, theta: float, weights: Tensor, reference: Tensor | None = None) -> None: if theta < 0.0: raise ValueError(f"Parameter `theta` should be non-negative. Found `theta = {theta}`.") super().__init__() self.theta = theta - self.weights = weights - self.reference = reference + self.register_buffer("weights", weights, persistent=False) + self.register_buffer("reference", reference, persistent=False) def forward(self, values: Tensor, /) -> Tensor: check_same_shape_as_values(self.weights, values, "weights") diff --git a/src/torchjd/scalarization/_stch.py b/src/torchjd/scalarization/_stch.py index 25ec95bc..019dd190 100644 --- a/src/torchjd/scalarization/_stch.py +++ b/src/torchjd/scalarization/_stch.py @@ -38,6 +38,9 @@ class STCH(Scalarizer): is applied. If provided, it must have the same shape as the values passed at call time. """ + weights: Tensor | None + reference: Tensor | None + def __init__( self, mu: float, @@ -49,8 +52,8 @@ def __init__( super().__init__() self.mu = mu - self.weights = weights - self.reference = reference + self.register_buffer("weights", weights, persistent=False) + self.register_buffer("reference", reference, persistent=False) def forward(self, values: Tensor, /) -> Tensor: check_same_shape_as_values(self.weights, values, "weights") diff --git a/tests/unit/scalarization/test_constant.py b/tests/unit/scalarization/test_constant.py index 126b91c0..1ae6f673 100644 --- a/tests/unit/scalarization/test_constant.py +++ b/tests/unit/scalarization/test_constant.py @@ -61,3 +61,11 @@ def test_representations() -> None: def test_str_with_non_vector_weights() -> None: assert str(Constant(weights=ones_((3, 4)))) == "Constant(weights of shape (3, 4))" assert str(Constant(weights=ones_(()))) == "Constant(weights of shape ())" + + +def test_to_moves_weights() -> None: + s = Constant(weights=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64) + + assert s.weights.dtype == torch.float64 + assert any(buffer is s.weights for buffer in s.buffers()) + assert len(s.state_dict()) == 0 diff --git a/tests/unit/scalarization/test_cosmos.py b/tests/unit/scalarization/test_cosmos.py index d98369d1..399aa8f3 100644 --- a/tests/unit/scalarization/test_cosmos.py +++ b/tests/unit/scalarization/test_cosmos.py @@ -86,3 +86,11 @@ def test_representations() -> None: s = COSMOS(lambda_=0.5, weights=torch.tensor([0.5, 0.5])) assert repr(s) == "COSMOS(lambda_=0.5, weights=tensor([0.5000, 0.5000]))" assert str(s) == "COSMOS" + + +def test_to_moves_weights() -> None: + s = COSMOS(lambda_=0.5, weights=tensor_([0.5, 0.5], dtype=torch.float32)).to(torch.float64) + + assert s.weights.dtype == torch.float64 + assert any(buffer is s.weights for buffer in s.buffers()) + assert len(s.state_dict()) == 0 diff --git a/tests/unit/scalarization/test_pbi.py b/tests/unit/scalarization/test_pbi.py index f7187d27..4d3e2d1d 100644 --- a/tests/unit/scalarization/test_pbi.py +++ b/tests/unit/scalarization/test_pbi.py @@ -103,3 +103,19 @@ def test_representations() -> None: s = PBI(theta=5.0, weights=torch.tensor([0.5, 0.5])) assert repr(s) == "PBI(theta=5.0, weights=tensor([0.5000, 0.5000]), reference=None)" assert str(s) == "PBI" + + +def test_to_moves_weights_and_reference() -> None: + s = PBI( + theta=5.0, + weights=tensor_([0.5, 0.5], dtype=torch.float32), + reference=tensor_([0.0, 1.0], dtype=torch.float32), + ).to(torch.float64) + + assert s.reference is not None + assert s.weights.dtype == torch.float64 + assert s.reference.dtype == torch.float64 + buffers = list(s.buffers()) + assert any(buffer is s.weights for buffer in buffers) + assert any(buffer is s.reference for buffer in buffers) + assert len(s.state_dict()) == 0 diff --git a/tests/unit/scalarization/test_stch.py b/tests/unit/scalarization/test_stch.py index 659a6a3b..2bdf7c37 100644 --- a/tests/unit/scalarization/test_stch.py +++ b/tests/unit/scalarization/test_stch.py @@ -80,3 +80,20 @@ def test_representations() -> None: s = STCH(mu=0.5) assert repr(s) == "STCH(mu=0.5, weights=None, reference=None)" assert str(s) == "STCH" + + +def test_to_moves_weights_and_reference() -> None: + s = STCH( + mu=1.0, + weights=tensor_([0.5, 0.5], dtype=torch.float32), + reference=tensor_([0.0, 1.0], dtype=torch.float32), + ).to(torch.float64) + + assert s.weights is not None + assert s.reference is not None + assert s.weights.dtype == torch.float64 + assert s.reference.dtype == torch.float64 + buffers = list(s.buffers()) + assert any(buffer is s.weights for buffer in buffers) + assert any(buffer is s.reference for buffer in buffers) + assert len(s.state_dict()) == 0