(self, net_type: str = 'alex', version: str = '0.1')
| 129 | version (str): the version of LPIPS. Default: 0.1. |
| 130 | """ |
| 131 | def __init__(self, net_type: str = 'alex', version: str = '0.1'): |
| 132 | |
| 133 | assert version in ['0.1'], 'v0.1 is only supported now' |
| 134 | |
| 135 | super(LPIPS, self).__init__() |
| 136 | |
| 137 | # pretrained network |
| 138 | self.net = get_network(net_type).to("cuda") |
| 139 | |
| 140 | # linear layers |
| 141 | self.lin = LinLayers(self.net.n_channels_list).to("cuda") |
| 142 | self.lin.load_state_dict(get_state_dict(net_type, version)) |
| 143 | |
| 144 | def forward(self, x: torch.Tensor, y: torch.Tensor): |
| 145 | feat_x, feat_y = self.net(x), self.net(y) |
nothing calls this directly
no test coverage detected