Rename attribute to feature

This commit is contained in:
henryruhs
2025-03-16 08:39:35 +01:00
parent 904a447e06
commit 94571c5676
10 changed files with 62 additions and 60 deletions
+6 -6
View File
@@ -5,7 +5,7 @@ from torch import Tensor, nn
from ..networks.aad import AAD
from ..networks.unet import UNet
from ..types import Attribute, Embedding
from ..types import Feature, Embedding
class Generator(nn.Module):
@@ -16,12 +16,12 @@ class Generator(nn.Module):
self.encoder.apply(init_weight)
self.generator.apply(init_weight)
def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Tuple[Attribute, ...]]:
target_attributes = self.encode_attributes(target_tensor)
output_tensor = self.generator(source_embedding, target_attributes)
return output_tensor, target_attributes
def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Tuple[Feature, ...]]:
target_features = self.encode_features(target_tensor)
output_tensor = self.generator(source_embedding, target_features)
return output_tensor, target_features
def encode_attributes(self, input_tensor : Tensor) -> Tuple[Attribute, ...]:
def encode_features(self, input_tensor : Tensor) -> Tuple[Feature, ...]:
return self.encoder(input_tensor)
+9 -9
View File
@@ -7,7 +7,7 @@ from torch import Tensor, nn
from torchvision import transforms
from ..helper import calc_embedding
from ..types import Attribute, EmbedderModule, FaceParserModule, GazerModule, Loss, Mask, MotionExtractorModule
from ..types import Feature, EmbedderModule, FaceParserModule, GazerModule, Loss, Mask, MotionExtractorModule
class DiscriminatorLoss(nn.Module):
@@ -49,22 +49,22 @@ class AdversarialLoss(nn.Module):
return adversarial_loss, weighted_adversarial_loss
class AttributeLoss(nn.Module):
class FeautureLoss(nn.Module):
def __init__(self, config_parser : ConfigParser) -> None:
super().__init__()
self.config_batch_size = config_parser.getint('training.loader', 'batch_size')
self.config_attribute_weight = config_parser.getfloat('training.losses', 'attribute_weight')
self.config_feature_weight = config_parser.getfloat('training.losses', 'feature_weight')
def forward(self, target_attributes : Tuple[Attribute, ...], output_attributes : Tuple[Attribute, ...]) -> Tuple[Loss, Loss]:
def forward(self, target_features : Tuple[Feature, ...], output_features : Tuple[Feature, ...]) -> Tuple[Loss, Loss]:
temp_tensors = []
for target_attribute, output_attribute in zip(target_attributes, output_attributes):
temp_tensor = torch.mean(torch.pow(output_attribute - target_attribute, 2).reshape(self.config_batch_size, -1), dim = 1).mean()
for target_feature, output_feature in zip(target_features, output_features):
temp_tensor = torch.mean(torch.pow(output_feature - target_feature, 2).reshape(self.config_batch_size, -1), dim = 1).mean()
temp_tensors.append(temp_tensor)
attribute_loss = torch.stack(temp_tensors).mean() * 0.5
weighted_attribute_loss = attribute_loss * self.config_attribute_weight
return attribute_loss, weighted_attribute_loss
feature_loss = torch.stack(temp_tensors).mean() * 0.5
weighted_feature_loss = feature_loss * self.config_feature_weight
return feature_loss, weighted_feature_loss
class ReconstructionLoss(nn.Module):