| 26 | # Learned perceptual metric |
| 27 | class PNetLin(nn.Module): |
| 28 | def __init__(self, pnet_type='vgg', pnet_rand=False, pnet_tune=False, use_dropout=True, spatial=False, version='0.1', lpips=True): |
| 29 | super(PNetLin, self).__init__() |
| 30 | |
| 31 | self.pnet_type = pnet_type |
| 32 | self.pnet_tune = pnet_tune |
| 33 | self.pnet_rand = pnet_rand |
| 34 | self.spatial = spatial |
| 35 | self.lpips = lpips |
| 36 | self.version = version |
| 37 | self.scaling_layer = ScalingLayer() |
| 38 | |
| 39 | if(self.pnet_type in ['vgg','vgg16']): |
| 40 | net_type = pn.vgg16 |
| 41 | self.chns = [64,128,256,512,512] |
| 42 | elif(self.pnet_type=='alex'): |
| 43 | net_type = pn.alexnet |
| 44 | self.chns = [64,192,384,256,256] |
| 45 | elif(self.pnet_type=='squeeze'): |
| 46 | net_type = pn.squeezenet |
| 47 | self.chns = [64,128,256,384,384,512,512] |
| 48 | self.L = len(self.chns) |
| 49 | |
| 50 | self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune) |
| 51 | |
| 52 | if(lpips): |
| 53 | self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout) |
| 54 | self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout) |
| 55 | self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout) |
| 56 | self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout) |
| 57 | self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout) |
| 58 | self.lins = [self.lin0,self.lin1,self.lin2,self.lin3,self.lin4] |
| 59 | if(self.pnet_type=='squeeze'): # 7 layers for squeezenet |
| 60 | self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout) |
| 61 | self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout) |
| 62 | self.lins+=[self.lin5,self.lin6] |
| 63 | |
| 64 | def forward(self, in0, in1, retPerLayer=False): |
| 65 | # v0.0 - original release had a bug, where input was not scaled |