|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
import sys, time, numpy, os, subprocess, pandas, tqdm
|
|
|
|
from loss_multi import lossAV, lossA, lossV
|
|
from model.loconet_encoder import locoencoder
|
|
|
|
import torch.distributed as dist
|
|
from xxlib.utils.distributed import all_gather, all_reduce
|
|
|
|
|
|
class Loconet(nn.Module):
|
|
|
|
def __init__(self, cfg):
|
|
super(Loconet, self).__init__()
|
|
self.cfg = cfg
|
|
self.model = locoencoder(cfg)
|
|
self.lossAV = lossAV()
|
|
self.lossA = lossA()
|
|
self.lossV = lossV()
|
|
|
|
def forward(self, audioFeature, visualFeature, labels, masks):
|
|
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)
|
|
|
|
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
|
|
|
|
num_frames = masks.sum()
|
|
return nloss, prec, num_frames
|
|
|
|
|
|
class loconet(nn.Module):
|
|
|
|
def __init__(self, cfg, rank=None, device=None):
|
|
super(loconet, self).__init__()
|
|
self.cfg = cfg
|
|
self.rank = rank
|
|
if rank != None:
|
|
self.rank = rank
|
|
self.device = device
|
|
|
|
self.model = Loconet(cfg).to(device)
|
|
self.model = nn.SyncBatchNorm.convert_sync_batchnorm(self.model)
|
|
self.model = nn.parallel.DistributedDataParallel(self.model,
|
|
device_ids=[rank],
|
|
output_device=rank,
|
|
find_unused_parameters=False)
|
|
self.optim = torch.optim.Adam(self.model.parameters(), lr=self.cfg.SOLVER.BASE_LR)
|
|
self.scheduler = torch.optim.lr_scheduler.StepLR(self.optim,
|
|
step_size=1,
|
|
gamma=self.cfg.SOLVER.SCHEDULER.GAMMA)
|
|
else:
|
|
self.model = locoencoder(cfg).cuda()
|
|
self.lossAV = lossAV().cuda()
|
|
self.lossA = lossA().cuda()
|
|
self.lossV = lossV().cuda()
|
|
|
|
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, epoch, loader):
|
|
self.model.train()
|
|
self.scheduler.step(epoch - 1)
|
|
index, top1, loss = 0, 0, 0
|
|
lr = self.optim.param_groups[0]['lr']
|
|
loader.sampler.set_epoch(epoch)
|
|
device = self.device
|
|
|
|
pbar = enumerate(loader, start=1)
|
|
if self.rank == 0:
|
|
pbar = tqdm.tqdm(pbar, total=loader.__len__())
|
|
|
|
for num, (audioFeature, visualFeature, labels, masks) in pbar:
|
|
|
|
audioFeature = audioFeature.to(device)
|
|
visualFeature = visualFeature.to(device)
|
|
labels = labels.to(device)
|
|
masks = masks.to(device)
|
|
nloss, prec, num_frames = self.model(
|
|
audioFeature,
|
|
visualFeature,
|
|
labels,
|
|
masks,
|
|
)
|
|
|
|
self.optim.zero_grad()
|
|
nloss.backward()
|
|
self.optim.step()
|
|
|
|
[nloss, prec, num_frames] = all_reduce([nloss, prec, num_frames], average=False)
|
|
top1 += prec.detach().cpu().numpy()
|
|
loss += nloss.detach().cpu().numpy()
|
|
index += int(num_frames.detach().cpu().item())
|
|
if self.rank == 0:
|
|
pbar.set_postfix(
|
|
dict(epoch=epoch,
|
|
lr=lr,
|
|
loss=loss / (num * self.cfg.NUM_GPUS),
|
|
acc=(top1 / index)))
|
|
dist.barrier()
|
|
return loss / num, lr
|
|
|
|
def evaluate_network(self, epoch, loader):
|
|
self.eval()
|
|
predScores = []
|
|
evalCsvSave = os.path.join(self.cfg.WORKSPACE, "{}_res.csv".format(epoch))
|
|
evalOrig = self.cfg.evalOrig
|
|
for audioFeature, visualFeature, labels, masks in tqdm.tqdm(loader):
|
|
with torch.no_grad():
|
|
audioFeature = audioFeature.cuda()
|
|
visualFeature = visualFeature.cuda()
|
|
labels = labels.cuda()
|
|
masks = masks.cuda()
|
|
b, s, t = visualFeature.shape[0], visualFeature.shape[1], visualFeature.shape[2]
|
|
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)
|
|
labels = labels.reshape((-1))
|
|
masks = masks.reshape((-1))
|
|
outsAV = outsAV.view(b, s, t, -1)[:, 0, :, :].view(b * t, -1)
|
|
labels = labels.view(b, s, t)[:, 0, :].view(b * t).cuda()
|
|
masks = masks.view(b, s, t)[:, 0, :].view(b * t)
|
|
_, predScore, _, _ = self.lossAV.forward(outsAV, labels, masks)
|
|
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, map_location='cpu')
|
|
if self.rank != None:
|
|
info = self.load_state_dict(loadedState)
|
|
else:
|
|
new_state = {}
|
|
|
|
for k, v in loadedState.items():
|
|
new_state[k.replace("model.module.", "")] = v
|
|
info = self.load_state_dict(new_state, strict=False)
|
|
print(info)
|
|
|
