| 26 | |
| 27 | # Learned perceptual metric |
| 28 | class PNetLin(nn.Module): |
| 29 | def __init__(self, pnet_type='vgg', pnet_rand=False, pnet_tune=False, use_dropout=True, spatial=False, version='0.1', lpips=True): |
| 30 | super(PNetLin, self).__init__() |
| 31 | |
| 32 | self.pnet_type = pnet_type |
| 33 | self.pnet_tune = pnet_tune |
| 34 | self.pnet_rand = pnet_rand |
| 35 | self.spatial = spatial |
| 36 | self.lpips = lpips |
| 37 | self.version = version |
| 38 | self.scaling_layer = ScalingLayer() |
| 39 | |
| 40 | if(self.pnet_type in ['vgg','vgg16']): |
| 41 | net_type = pn.vgg16 |
| 42 | self.chns = [64,128,256,512,512] |
| 43 | elif(self.pnet_type=='alex'): |
| 44 | net_type = pn.alexnet |
| 45 | self.chns = [64,192,384,256,256] |
| 46 | elif(self.pnet_type=='squeeze'): |
| 47 | net_type = pn.squeezenet |
| 48 | self.chns = [64,128,256,384,384,512,512] |
| 49 | self.L = len(self.chns) |
| 50 | |
| 51 | self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune) |
| 52 | |
| 53 | if(lpips): |
| 54 | self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout) |
| 55 | self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout) |
| 56 | self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout) |
| 57 | self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout) |
| 58 | self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout) |
| 59 | self.lins = [self.lin0,self.lin1,self.lin2,self.lin3,self.lin4] |
| 60 | if(self.pnet_type=='squeeze'): # 7 layers for squeezenet |
| 61 | self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout) |
| 62 | self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout) |
| 63 | self.lins+=[self.lin5,self.lin6] |
| 64 | |
| 65 | def forward(self, in0, in1, retPerLayer=False): |
| 66 | # v0.0 - original release had a bug, where input was not scaled |
| 67 | in0_input, in1_input = (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version=='0.1' else (in0, in1) |
| 68 | outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input) |
| 69 | feats0, feats1, diffs = {}, {}, {} |
| 70 | |
| 71 | for kk in range(self.L): |
| 72 | feats0[kk], feats1[kk] = util.normalize_tensor(outs0[kk]), util.normalize_tensor(outs1[kk]) |
| 73 | diffs[kk] = (feats0[kk]-feats1[kk])**2 |
| 74 | |
| 75 | if(self.lpips): |
| 76 | if(self.spatial): |
| 77 | res = [upsample(self.lins[kk].model(diffs[kk]), out_H=in0.shape[2]) for kk in range(self.L)] |
| 78 | else: |
| 79 | res = [spatial_average(self.lins[kk].model(diffs[kk]), keepdim=True) for kk in range(self.L)] |
| 80 | else: |
| 81 | if(self.spatial): |
| 82 | res = [upsample(diffs[kk].sum(dim=1,keepdim=True), out_H=in0.shape[2]) for kk in range(self.L)] |
| 83 | else: |
| 84 | res = [spatial_average(diffs[kk].sum(dim=1,keepdim=True), keepdim=True) for kk in range(self.L)] |
| 85 |
nothing calls this directly
no outgoing calls
no test coverage detected