34 lines
793 B
Python
34 lines
793 B
Python
|
import torch
|
||
|
|
|
||
|
|
from network.SbS import SbS
|
||
|
|
|
||
|
|
|
||
|
|
class SbSReconstruction(torch.nn.Module):
|
||
|
|
|
||
|
|
_the_sbs_layer: SbS
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
the_sbs_layer: SbS,
|
||
|
|
) -> None:
|
||
|
|
super().__init__()
|
||
|
|
|
||
|
|
self._the_sbs_layer = the_sbs_layer
|
||
|
|
self.device = self._the_sbs_layer.device
|
||
|
|
self.default_dtype = self._the_sbs_layer.default_dtype
|
||
|
|
|
||
|
|
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
||
|
|
|
||
|
|
assert self._the_sbs_layer._weights_exists is True
|
||
|
|
|
||
|
|
input_norm = input / input.sum(dim=1, keepdim=True)
|
||
|
|
|
||
|
|
output = (
|
||
|
|
self._the_sbs_layer._weights.data.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
|
||
|
|
* input_norm.unsqueeze(1)
|
||
|
|
).sum(dim=2)
|
||
|
|
|
||
|
|
output /= output.sum(dim=1, keepdim=True)
|
||
|
|
|
||
|
|
return output
|