| | from config_loconet import LoCoNetConfig |
| | from transformers import PreTrainedModel |
| | from loconet_encoder import locoencoder |
| | from loss_multi import lossAV, lossA, lossV |
| |
|
| |
|
| | class loconet(PreTrainedModel): |
| | config_class = LoCoNetConfig |
| |
|
| | def __init__(self, config): |
| | super().__init__(config) |
| |
|
| | self.model = locoencoder(config) |
| | self.lossAV = lossAV() |
| | self.lossA = lossA() |
| | self.lossV = lossV() |
| |
|
| | def forward(self, audioFeature, visualFeature, masks, labels=None): |
| | b, s, t = visualFeature.shape[:3] |
| | visualFeature = visualFeature.view(b * s, *visualFeature.shape[2:]) |
| | labels = labels.view(b * s, *labels.shape[2:]) |
| | masks = masks.view(b * s, *masks.shape[2:]) |
| |
|
| | audioEmbed = self.model.forward_audio_frontend(audioFeature) |
| | visualEmbed = self.model.forward_visual_frontend(visualFeature) |
| | audioEmbed = audioEmbed.repeat(s, 1, 1) |
| |
|
| | audioEmbed, visualEmbed = self.model.forward_cross_attention(audioEmbed, visualEmbed) |
| | outsAV = self.model.forward_audio_visual_backend(audioEmbed, visualEmbed, b, s) |
| | outsA = self.model.forward_audio_backend(audioEmbed) |
| | outsV = self.model.forward_visual_backend(visualEmbed) |
| | num_frames = masks.sum() |
| |
|
| | if labels is not None: |
| |
|
| | labels = labels.reshape((-1)) |
| | masks = masks.reshape((-1)) |
| | nlossAV, _, _, prec = self.lossAV.forward(outsAV, labels, masks) |
| | nlossA = self.lossA.forward(outsA, labels, masks) |
| | nlossV = self.lossV.forward(outsV, labels, masks) |
| |
|
| | nloss = nlossAV + 0.4 * nlossA + 0.4 * nlossV |
| |
|
| | return {"loss": nloss, "logits": outsAV} |
| |
|
| | else: |
| |
|
| | return {"logits": outsAV} |
| |
|
