class FeatureLoss(nn.Module):
def __init__(self, m_feat, layer_ids, layer_wgts):
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))]
When the FeatureLoss class is created with the vgg model and three layer ids, self.loss_features is just :[ReLU(inplace=True), ReLU(inplace=True), ReLU(inplace=True)]
When self.loss_features is passed to hook_outputs, how does it know which RELU layers from the VGG model to create hooks for? We are not passing the layer ids to the hook_outputs function.