| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
| import sys, time, numpy, os, subprocess, pandas, tqdm
|
|
|
| from loss import lossAV, lossA, lossV
|
| from model.talkNetModel import talkNetModel
|
|
|
| class talkNet(nn.Module):
|
| def __init__(self, lr = 0.0001, lrDecay = 0.95, **kwargs):
|
| super(talkNet, self).__init__()
|
| self.model = talkNetModel().cuda()
|
| self.lossAV = lossAV().cuda()
|
| self.lossA = lossA().cuda()
|
| self.lossV = lossV().cuda()
|
| self.optim = torch.optim.Adam(self.parameters(), lr = lr)
|
| self.scheduler = torch.optim.lr_scheduler.StepLR(self.optim, step_size = 1, gamma=lrDecay)
|
| print(time.strftime("%m-%d %H:%M:%S") + " Model para number = %.2f"%(sum(param.numel() for param in self.model.parameters()) / 1024 / 1024))
|
|
|
| def train_network(self, loader, epoch, **kwargs):
|
| self.train()
|
| self.scheduler.step(epoch - 1)
|
| index, top1, loss = 0, 0, 0
|
| lr = self.optim.param_groups[0]['lr']
|
| for num, (audioFeature, visualFeature, labels) in enumerate(loader, start=1):
|
| self.zero_grad()
|
| audioEmbed = self.model.forward_audio_frontend(audioFeature[0].cuda())
|
| visualEmbed = self.model.forward_visual_frontend(visualFeature[0].cuda())
|
| audioEmbed, visualEmbed = self.model.forward_cross_attention(audioEmbed, visualEmbed)
|
| outsAV= self.model.forward_audio_visual_backend(audioEmbed, visualEmbed)
|
| outsA = self.model.forward_audio_backend(audioEmbed)
|
| outsV = self.model.forward_visual_backend(visualEmbed)
|
| labels = labels[0].reshape((-1)).cuda()
|
| nlossAV, _, _, prec = self.lossAV.forward(outsAV, labels)
|
| nlossA = self.lossA.forward(outsA, labels)
|
| nlossV = self.lossV.forward(outsV, labels)
|
| nloss = nlossAV + 0.4 * nlossA + 0.4 * nlossV
|
| loss += nloss.detach().cpu().numpy()
|
| top1 += prec
|
| nloss.backward()
|
| self.optim.step()
|
| index += len(labels)
|
| sys.stderr.write(time.strftime("%m-%d %H:%M:%S") + \
|
| " [%2d] Lr: %5f, Training: %.2f%%, " %(epoch, lr, 100 * (num / loader.__len__())) + \
|
| " Loss: %.5f, ACC: %2.2f%% \r" %(loss/(num), 100 * (top1/index)))
|
| sys.stderr.flush()
|
| sys.stdout.write("\n")
|
| return loss/num, lr
|
|
|
| def evaluate_network(self, loader, evalCsvSave, evalOrig, **kwargs):
|
| self.eval()
|
| predScores = []
|
| for audioFeature, visualFeature, labels in tqdm.tqdm(loader):
|
| with torch.no_grad():
|
| audioEmbed = self.model.forward_audio_frontend(audioFeature[0].cuda())
|
| visualEmbed = self.model.forward_visual_frontend(visualFeature[0].cuda())
|
| audioEmbed, visualEmbed = self.model.forward_cross_attention(audioEmbed, visualEmbed)
|
| outsAV= self.model.forward_audio_visual_backend(audioEmbed, visualEmbed)
|
| labels = labels[0].reshape((-1)).cuda()
|
| _, predScore, _, _ = self.lossAV.forward(outsAV, labels)
|
| predScore = predScore[:,1].detach().cpu().numpy()
|
| predScores.extend(predScore)
|
| evalLines = open(evalOrig).read().splitlines()[1:]
|
| labels = []
|
| labels = pandas.Series( ['SPEAKING_AUDIBLE' for line in evalLines])
|
| scores = pandas.Series(predScores)
|
| evalRes = pandas.read_csv(evalOrig)
|
| evalRes['score'] = scores
|
| evalRes['label'] = labels
|
| evalRes.drop(['label_id'], axis=1,inplace=True)
|
| evalRes.drop(['instance_id'], axis=1,inplace=True)
|
| evalRes.to_csv(evalCsvSave, index=False)
|
| cmd = "python -O utils/get_ava_active_speaker_performance.py -g %s -p %s "%(evalOrig, evalCsvSave)
|
| mAP = float(str(subprocess.run(cmd, shell=True, capture_output =True).stdout).split(' ')[2][:5])
|
| return mAP
|
|
|
| def saveParameters(self, path):
|
| torch.save(self.state_dict(), path)
|
|
|
| def loadParameters(self, path):
|
| selfState = self.state_dict()
|
| loadedState = torch.load(path)
|
| for name, param in loadedState.items():
|
| origName = name;
|
| if name not in selfState:
|
| name = name.replace("module.", "")
|
| if name not in selfState:
|
| print("%s is not in the model."%origName)
|
| continue
|
| if selfState[name].size() != loadedState[origName].size():
|
| sys.stderr.write("Wrong parameter length: %s, model: %s, loaded: %s"%(origName, selfState[name].size(), loadedState[origName].size()))
|
| continue
|
| selfState[name].copy_(param)
|
|
|
