def create_loss(device_type='cuda'):
if device_type == 'cuda':
vgg_m = vgg16_bn(True).features.cuda().eval()
else:
vgg_m = vgg16_bn(True).features.eval()
requires_grad(vgg_m, False)
blocks = [
i - 1 for i, o in enumerate(children(vgg_m))
if isinstance(o, nn.MaxPool2d)
]
blocks, [vgg_m[i] for i in blocks]
feat_loss = FeatureLoss(vgg_m, blocks[2:5], [5, 15, 2])
return feat_loss
def __init__(self, layer_wgts=[20, 70, 10]):
super().__init__()
self.m_feat = models.vgg16_bn(True).features.cuda().eval()
requires_grad(self.m_feat, False)
blocks = [
i - 1 for i, o in enumerate(children(self.m_feat))
if isinstance(o, nn.MaxPool2d)
]
layer_ids = blocks[2:5]
self.loss_features = [self.m_feat[i] for i in layer_ids]
self.hooks = hook_outputs(self.loss_features, detach=False)
self.wgts = layer_wgts
self.metric_names = ['pixel'] + \
[f'feat_{i}' for i in range(len(layer_ids))]
self.base_loss = F.l1_loss
conv_new = nn.Conv2d(31, 64, kernel_size=3, stride=1, padding=1)
weight_old = conv1.weight
weight_new = conv_new.weight
for i in range(10):
weight_new[:, i:i + 3] = weight_old.data.clone()
weight_new[:, 30] = weight_old.data[:, 0].clone()
conv_new.weight = nn.Parameter(weight_new)
setattr(vgg_m, '0', conv_new)
vgg_m = vgg_m.cuda().eval()
requires_grad(vgg_m, False)
blocks = [
i - 1 for i, o in enumerate(children(vgg_m))
if isinstance(o, nn.MaxPool2d)
]
blocks, [vgg_m[i] for i in blocks]
class FeatureLoss2(nn.Module):
def __init__(self, m_feat, layer_ids, layer_wgts, mrse=True):
super().__init__()
self.m_feat = m_feat
self.loss_features = [self.m_feat[i] for i in layer_ids]
self.hooks = hook_outputs(self.loss_features, detach=False)
self.wgts = layer_wgts
self.metric_names = [
'pixel',
] + [f'feat_{i}' for i in range(len(layer_ids))
conv1 = getattr(vgg_m, '0')
conv_new = nn.Conv2d(31, 64, kernel_size=3, stride=1, padding=1)
weight_old = conv1.weight
weight_new = conv_new.weight
for i in range(10):
weight_new[:, i:i+3] = weight_old.data.clone()
weight_new[:, 30] = weight_old.data[:, 0].clone()
conv_new.weight = nn.Parameter(weight_new)
setattr(vgg_m, '0', conv_new)
vgg_m = vgg_m.cuda().eval()
requires_grad(vgg_m, False)
blocks = [i-1 for i,o in enumerate(children(vgg_m)) if isinstance(o,nn.MaxPool2d)]
blocks, [vgg_m[i] for i in blocks]
class FeatureLoss2(nn.Module):
def __init__(self, m_feat, layer_ids, layer_wgts,mrse=True):
super().__init__()
self.m_feat = m_feat
self.loss_features = [self.m_feat[i] for i in layer_ids]
self.hooks = hook_outputs(self.loss_features, detach=False)
self.wgts = layer_wgts
self.metric_names = ['pixel',] + [f'feat_{i}' for i in range(len(layer_ids))
] + [f'gram_{i}' for i in range(len(layer_ids))]
self.mrae = MRAELoss()
self.mrse = MRSELoss()
self.mse = MSELossFlat()
self.mrse_switch = mrse