TorchJD-VAE/custom_aggregation.py at master · mattbuot/TorchJD-VAE · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
# from typing import Literal

# import torch
# from torch import Tensor
# from torchjd.aggregation._pref_vector_utils import (
#     pref_vector_to_str_suffix,
#     pref_vector_to_weighting,
# )
# from torchjd.aggregation.bases import _WeightedAggregator, _Weighting
# from torchjd.aggregation.mean import _MeanWeighting
# from torchjd.aggregation.upgrad import _UPGradWrapper


# class PairwiseUpgradAggregator(_WeightedAggregator):
#     def __init__(
#         self,
#         pref_vector: Tensor | None = None,
#         norm_eps: float = 0.0001,
#         reg_eps: float = 0.0001,
#         solver: Literal["quadprog"] = "quadprog",
#         final_aggregation: Literal["mean", "upgrad"] = "mean",
#     ):
#         weighting = pref_vector_to_weighting(pref_vector, default=_MeanWeighting())
#         self._pref_vector = pref_vector
#         self._final_aggregation = final_aggregation

#         super().__init__(
#             weighting=PairwiseUpgradAggregatorWrapper(
#                 weighting=weighting, norm_eps=norm_eps, reg_eps=reg_eps, solver=solver, final_aggregation=final_aggregation,
#             )
#         )

#     def __repr__(self) -> str:
#         return (
#             f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, norm_eps="
#             f"{self.weighting.norm_eps}, reg_eps={self.weighting.reg_eps}, "
#             f"solver={repr(self.weighting.solver)})"
#         )

#     def __str__(self) -> str:
#         return f"PairwiseUpgrad_{self._final_aggregation}{pref_vector_to_str_suffix(self._pref_vector)}"


# class PairwiseUpgradAggregatorWrapper(_Weighting):

#     def __init__(
#         self,
#         weighting: _Weighting,
#         norm_eps: float,
#         reg_eps: float,
#         solver: Literal["quadprog"],
#         final_aggregation: Literal["mean", "upgrad"],
#     ):
#         super().__init__()
#         self._upgrad_wrapper = _UPGradWrapper(
#             weighting=weighting, norm_eps=norm_eps, reg_eps=reg_eps, solver=solver
#         )
#         self.final_aggregation = final_aggregation

#     def forward(self, matrix: Tensor) -> Tensor:

#         pairwise_aggregations = torch.stack([
#             self._upgrad_wrapper(matrix[i:i + 2]) for i in range(0, matrix.size(0), 2)
#         ]).flatten()
#         return pairwise_aggregations / matrix.size(0)


#         # .view(1, -1)
#         # paired_matrix = torch.mul(pairwise_aggregations, matrix)

#         # if self.final_aggregation == "upgrad":
#         #     return self._upgrad_wrapper(paired_matrix)
#         # elif self.final_aggregation == "mean":
#         #     return paired_matrix.mean(dim=0)
#         # else:
#         #     raise ValueError(f"Unknown final aggregation method: {self.final_aggregation}")