#! /usr/bin/python
# -*- encoding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.commons.hparams import hparams
import numpy as np
import math
class LossScale(nn.Module):
def __init__(self, init_w=10.0, init_b=-5.0):
super(LossScale, self).__init__()
self.wC = nn.Parameter(torch.tensor(init_w))
self.bC = nn.Parameter(torch.tensor(init_b))
class CLIPLoss(nn.Module):
def __init__(self,):
super().__init__()
def forward(self, audio_features, motion_features, logit_scale, clip_mask=None):
logits_per_audio = logit_scale * audio_features @ motion_features.T # [b,c]
logits_per_motion = logit_scale * motion_features @ audio_features.T # [b,c]
if clip_mask is not None:
logits_per_audio += clip_mask
logits_per_motion += clip_mask
labels = torch.arange(logits_per_motion.shape[0]).to(logits_per_motion.device)
motion_loss = F.cross_entropy(logits_per_motion, labels)
audio_loss = F.cross_entropy(logits_per_audio, labels)
clip_loss = (motion_loss + audio_loss) / 2
ret = {
"audio_loss": audio_loss,
"motion_loss": motion_loss,
"clip_loss": clip_loss
}
return ret
def accuracy(output, target, topk=(1,)):
"""Computes the precision@k for the specified values of k"""
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.reshape(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].reshape(-1).float()
res.append(correct_k)
return res
class LossScale(nn.Module):
def __init__(self, init_w=10.0, init_b=-5.0):
super(LossScale, self).__init__()
self.wC = nn.Parameter(torch.tensor(init_w))
self.bC = nn.Parameter(torch.tensor(init_b))
class SyncNetModel(nn.Module):
def __init__(self, auddim=1024, lipdim=20*3, nOut = 1024, stride=1):
super(SyncNetModel, self).__init__()
self.loss_scale = LossScale()
self.criterion = torch.nn.CrossEntropyLoss(reduction='none')
self.clip_loss_fn = CLIPLoss()
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
self.logit_scale_max = math.log(1. / 0.01)
self.netcnnaud = nn.Sequential(
nn.Conv1d(auddim, 512, kernel_size=3, stride=1, padding=1),
nn.BatchNorm1d(512),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=3, stride=1),
nn.Conv1d(512, 512, kernel_size=3, stride=1, padding=1),
nn.BatchNorm1d(512),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=3, stride=1),
nn.Conv1d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm1d(512),
nn.ReLU(inplace=True),
nn.Conv1d(512, 256, kernel_size=3, padding=1),
nn.BatchNorm1d(256),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=3, stride=1),
nn.Conv1d(256, 256, kernel_size=3, padding=1),
nn.BatchNorm1d(256),
nn.ReLU(inplace=True),
nn.Conv1d(256, 512, kernel_size=3, padding=1, stride=(stride)),
nn.BatchNorm1d(512),
nn.ReLU(),
nn.MaxPool1d(kernel_size=3, stride=1),
nn.Conv1d(512, 512, kernel_size=2),
nn.BatchNorm1d(512),
nn.ReLU(),
nn.Conv1d(512, 512, kernel_size=1),
nn.BatchNorm1d(512),
nn.ReLU(),
nn.Conv1d(512, nOut, kernel_size=1),
)
self.netcnnlip = nn.Sequential(
nn.Conv1d(lipdim, 512, kernel_size=3, stride=1, padding=1),
nn.BatchNorm1d(512),
nn.ReLU(inplace=True),
nn.Conv1d(512, 512, kernel_size=3, stride=1, padding=1),
nn.BatchNorm1d(512),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=3, stride=1),
nn.Conv1d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm1d(512),
nn.ReLU(inplace=True),
nn.Conv1d(512, 256, kernel_size=3, padding=1),
nn.BatchNorm1d(256),
nn.ReLU(inplace=True),
nn.Conv1d(256, 256, kernel_size=3, padding=1),
nn.BatchNorm1d(256),
nn.ReLU(inplace=True),
nn.Conv1d(256, 512, kernel_size=(3), padding=1, stride=(stride)),
nn.BatchNorm1d(512),
nn.ReLU(),
nn.MaxPool1d(kernel_size=3, stride=1),
nn.Conv1d(512, 512, kernel_size=1),
nn.BatchNorm1d(512),
nn.ReLU(),
nn.Conv1d(512, nOut, kernel_size=1),
)
def _forward_aud(self, x):
# bct
out = self.netcnnaud(x); # N x ch x 24 x M
return out
def _forward_vid(self, x):
# bct
out = self.netcnnlip(x);
return out
def forward(self, hubert, mouth_lm):
# hubert := (B, T=100, C=1024)
# mouth_lm3d := (B, T=50, C=60)
# out: [B, T=50, C=1024]
hubert = hubert.transpose(1,2)
mouth_lm = mouth_lm.transpose(1,2)
mouth_embedding = self._forward_vid(mouth_lm)
audio_embedding = self._forward_aud(hubert)
audio_embedding = audio_embedding.transpose(1,2)
mouth_embedding = mouth_embedding.transpose(1,2)
if hparams.get('normalize_embedding', False): # similar loss, no effects
audio_embedding = F.normalize(audio_embedding, p=2, dim=-1)
mouth_embedding = F.normalize(mouth_embedding, p=2, dim=-1)
return audio_embedding.squeeze(1), mouth_embedding.squeeze(1)
def _compute_sync_loss_batch(self, out_a, out_v, ymask=None):
b, t, c = out_v.shape
label = torch.arange(t).to(out_v.device)[None].repeat(b, 1)
output = F.cosine_similarity(
out_v[:, :, None], out_a[:, None, :], dim=-1) * self.loss_scale.wC + self.loss_scale.bC
loss = self.criterion(output, label).mean()
return loss
def _compute_sync_loss(self, out_a, out_v, ymask=None):
# b,t,c
b, t, c = out_v.shape
out_v = out_v.transpose(1,2)
out_a = out_a.transpose(1,2)
label = torch.arange(t).to(out_v.device)
nloss = 0
prec1 = 0
if ymask is not None:
total_num = ymask.sum()
else:
total_num = b*t
for i in range(0, b):
ft_v = out_v[[i],:,:].transpose(2,0)
ft_a = out_a[[i],:,:].transpose(2,0)
output = F.cosine_similarity(ft_v, ft_a.transpose(0,2)) * self.loss_scale.wC + self.loss_scale.bC
loss = self.criterion(output, label)
if ymask is not None:
loss = loss * ymask[i]
nloss += loss.sum()
nloss = nloss / total_num
return nloss
def compute_sync_loss(self,out_a, out_v, ymask=None, batch_mode=False):
if batch_mode:
return self._compute_sync_loss_batch(out_a, out_v)
else:
return self._compute_sync_loss(out_a, out_v)
def compute_sync_score_for_infer(self, out_a, out_v, ymask=None):
# b,t,c
b, t, c = out_v.shape
out_v = out_v.transpose(1,2)
out_a = out_a.transpose(1,2)
label = torch.arange(t).to(out_v.device)
nloss = 0
prec1 = 0
if ymask is not None:
total_num = ymask.sum()
else:
total_num = b*t
for i in range(0, b):
ft_v = out_v[[i],:,:].transpose(2,0)
ft_a = out_a[[i],:,:].transpose(2,0)
output = F.cosine_similarity(ft_v, ft_a.transpose(0,2)) * self.loss_scale.wC + self.loss_scale.bC
loss = self.criterion(output, label)
if ymask is not None:
loss = loss * ymask[i]
nloss += loss.sum()
nloss = nloss / total_num
return nloss
def cal_clip_loss(self, audio_embedding, mouth_embedding):
logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
clip_ret = self.clip_loss_fn(audio_embedding, mouth_embedding, logit_scale)
loss = clip_ret['clip_loss']
return loss
if __name__ == '__main__':
syncnet = SyncNetModel()
aud = torch.randn([2, 10, 1024])
vid = torch.randn([2, 5, 60])
aud_feat, vid_feat = syncnet.forward(aud, vid)
print(aud_feat.shape)
print(vid_feat.shape)