def job(item):
fn = item
outpath = os.path.join(fn, 'flow.mp4')
if not os.path.exists(outpath):
flows = torch.stack([
torch.from_numpy(read_flow(_))
for _ in glob(os.path.join(fn, '*.flo'))
])
flows = list(normalize_flows(flows))
flows = list(flows)
rgb_flows = [make_uint8(flow2rgb(_.numpy())) for _ in flows]
vid = ImageSequenceClip(rgb_flows, fps=8)
vid.write_videofile(outpath, fps=8, verbose=False, logger=None)
vid.close()
def test(loader, save_flag, epoch):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
accs = AverageMeter()
model.eval()
gc = LayerGradCam(model, model.layer4)
loss_fn = nn.CrossEntropyLoss()
end = time.time()
if save_flag:
results = [[
'y', 'y_hat_vec', 'y_hat', 'viz_fn', 'fn', 't_start', 't_end'
]]
with tqdm(loader, desc="Test batch iteration",
disable=args.local_rank > 0) as t:
for batch_idx, (xs, ys, (fns, t_starts, t_ends)) in enumerate(t):
data_time.update(time.time() - end)
xs = xs.to(device)
ys = ys.to(device)
y_hats = model(xs)
loss = loss_fn(y_hats, ys)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
losses.update(loss.item(), len(ys))
accs.update(accuracy(y_hats, ys)[0].item(), len(ys))
batch_time.update(time.time() - end)
end = time.time()
d = 0
if save_flag:
for x, y, y_hat, fn, t_s, t_e in zip(xs, ys,
F.softmax(y_hats, dim=1),
fns, t_starts, t_ends):
x = unnormalize(x.cpu()).permute(1, 2, 3, 0).numpy()
fn_ = fn
fn = '{0:02}_{1:010}.mp4'.format(
args.local_rank, batch_idx * args.batch_size + d)
results.append(
(y.item(), y_hat.tolist(), y_hat.argmax().item(), fn,
fn_, t_s.item(), t_e.item()))
clip = ImageSequenceClip(list(x),
fps=args.fps).fl_image(make_uint8)
clip.write_videofile(os.path.join(args.save_path, 'input',
fn),
logger=None)
clip.close()
d += 1
t.set_postfix(DataTime=data_time.avg,
BatchTime=batch_time.avg,
Loss=losses.avg,
Acc=accs.avg)
if save_flag == True:
with open(
os.path.join(
args.save_path,
'results_{0:06}_{1:03}.csv'.format(args.local_rank,
epoch)), 'w') as f:
wr = csv.writer(f)
wr.writerows(results)
return accs.avg, accs.count
def test(loader, save_flag, epoch):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
sup_losses = AverageMeter()
selfsup_acc = AverageMeter()
selfsup_losses = AverageMeter()
accs = AverageMeter()
if not args.save_attn:
model.eval()
else:
# ipdb.set_trace()
model.module.fc = selfsup_model[0]
gc = LayerGradCam(model, model.module.layer4)
loss_fn = nn.CrossEntropyLoss(ignore_index=-1)
if args.selfsup_loss == 'pred_middle':
selfsup_loss_fn = nn.MSELoss()
elif args.selfsup_loss == 'sort' or args.selfsup_loss == 'fps':
selfsup_loss_fn = loss_fn
elif args.selfsup_loss == 'ctc':
selfsup_loss_fn = ctc_loss
end = time.time()
if save_flag:
results = [['y', 'y_hat_vec', 'y_hat', 'viz_fn', 'fn', 't_start', 't_end']]
if args.flow_histogram:
results = [['x', 'y', 'y_hat_vec', 'y_hat', 'viz_fn', 'fn', 't_start', 't_end']]
featsarr = []
predsarr =[]
with tqdm(loader, desc="Test batch iteration", disable=args.local_rank > 0) as t:
for batch_idx, (xs, ys, (fns, t_starts, t_ends, selfsup_info, *_)) in enumerate(t):
data_time.update(time.time() - end)
xs = xs.to(device)
ys = ys.to(device)
if args.get_features:
_, feats = model(xs)
for feat, fn, t_start, t_end in zip(feats.detach().cpu(), fns, t_starts, t_ends):
featsarr.append((feat, fn, t_start, t_end))
continue
if args.save_preds:
_, feats = model(xs)
pred_fps = selfsup_model(feats).argmax(1)
for pred, fn, t_start, t_end in zip(pred_fps.detach().cpu(), fns, t_starts, t_ends):
predsarr.append((pred.item(), fn, t_start.item(), t_end.item()))
continue
if args.save_attn:
with torch.no_grad():
y_hats = model(xs)
if args.local_rank <= 0: ipdb.set_trace()
yh_argmax = y_hats.argmax(dim=1)
xs.requires_grad = True
fps_ys = torch.LongTensor([args.fps_list.index(_) for _ in selfsup_info]).to(device)
attr = gc.attribute(xs, yh_argmax)
up_attr = LayerAttribution.interpolate(attr, (16, 112, 112), interpolate_mode='trilinear').to(torch.float)
xs_ = torch.stack([unnormalize(x.cpu()) for x in xs])
acts = xs_.cpu() * up_attr.cpu()
acts = acts.cpu().detach().clamp(min=0)
for act, fn, t_s, t_e, yh, y in zip(acts, fns, t_starts, t_ends, yh_argmax.tolist(), fps_ys.tolist()):
# if args.local_rank <= 0: ipdb.set_trace()
save_image(act.permute(1, 0, 2, 3), os.path.join(args.save_path, 'input', f'{os.path.splitext(os.path.basename(fn))[0]}_{int(1000*t_s)}_{int(1000*t_e)}_pred{yh}_gt{y}.png'), normalize=True)
accs.update(accuracy(y_hats, fps_ys)[0].item(), len(fps_ys))
t.set_postfix(
Acc=accs.avg
)
continue
if args.selfsup_loss:
if args.selfsup_loss == 'pred_middle' or args.selfsup_loss == 'ctc':
_, prev_feats = model(xs[:, 0])
y_hats, mid_feats = model(xs[:, 1])
_, next_feats = model(xs[:, 2])
feats = torch.cat((prev_feats, next_feats), dim=1)
pred_mid_feats = selfsup_model(feats)
valid_pred_locs = (xs[:, 0].mean(dim=(1, 2, 3, 4)) > -0.999) & (
xs[:, 2].mean(dim=(1, 2, 3, 4)) > -0.999)
if args.selfsup_loss == 'pred_middle':
selfsup_loss = selfsup_loss_fn(pred_mid_feats[valid_pred_locs], mid_feats[valid_pred_locs])
elif args.selfsup_loss == 'ctc':
selfsup_loss = selfsup_loss_fn(pred_mid_feats[valid_pred_locs], mid_feats[valid_pred_locs],
feats[valid_pred_locs])
selfsup_len = valid_pred_locs.sum().item()
elif args.selfsup_loss == 'sort':
sort_ys = torch.zeros_like(ys)
valid_pred_locs = (xs[:, 0].mean(dim=(1, 2, 3, 4)) > -0.999) & (
xs[:, 2].mean(dim=(1, 2, 3, 4)) > -0.999)
for i in range(len(xs)):
p = torch.randperm(3)
xs[i] = xs[i][p]
s = ''.join(map(str, p.tolist()))
try:
sort_ys[i] = sort_y_vocab.index(s)
except:
sort_ys[i] = sort_y_vocab.index(s[::-1])
_, prev_feats = model(xs[:, 0])
y_hats, mid_feats = model(xs[:, 1]) # nonsense, can't co train with sort for now
_, next_feats = model(xs[:, 2])
feats = torch.stack((prev_feats, mid_feats, next_feats), dim=1)
pred_perms = selfsup_model(feats)
sort_ys[~valid_pred_locs] = -1
selfsup_loss = selfsup_loss_fn(pred_perms, sort_ys)
selfsup_len = valid_pred_locs.sum().item()
selfsup_acc.update(accuracy(pred_perms[valid_pred_locs], sort_ys[valid_pred_locs])[0].item(),
selfsup_len)
elif args.selfsup_loss == 'fps':
fps_ys = torch.LongTensor([args.fps_list.index(_) for _ in selfsup_info]).to(device)
y_hats, feats = model(xs)
pred_fps = selfsup_model(feats)
selfsup_loss = selfsup_loss_fn(pred_fps, fps_ys)
selfsup_len = len(ys)
selfsup_acc.update(accuracy(pred_fps, fps_ys)[0].item(), selfsup_len)
suploss = loss_fn(y_hats, ys)
loss = suploss + args.selfsup_lambda * selfsup_loss
else:
y_hats = model(xs)
suploss = loss_fn(y_hats, ys)
loss = suploss
# print(loss, y_hats, ys)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
losses.update(loss.item(), len(ys))
if args.selfsup_loss:
if (ys != -1).sum() > 0:
sup_losses.update(suploss.item(), (ys != -1).sum().item())
accs.update(accuracy(y_hats[ys != -1], ys[ys != -1])[0].item(), len(ys))
selfsup_losses.update(selfsup_loss.item(), selfsup_len)
else:
accs.update(accuracy(y_hats[ys != -1], ys[ys != -1])[0].item(), len(ys))
batch_time.update(time.time() - end)
end = time.time()
d = 0
if save_flag:
# TODO for self-supervised losses
for x, y, y_hat, fn, t_start, t_end in zip(xs, ys,
F.softmax(y_hats, dim=1),
fns, t_starts, t_ends):
fn_ = fn
fn = '{0:02}_{1:010}.mp4'.format(
args.local_rank, batch_idx * args.batch_size + d)
other = ()
if args.flow_histogram:
other = (x.tolist(),)
results.append(
(
*other, y.item(), y_hat.tolist(), y_hat.argmax().item(), fn, fn_, t_start.item(),
t_end.item()))
if args.save_test_vids:
x = unnormalize(x.cpu()).permute(1, 2, 3, 0).numpy()
tt = ImageSequenceClip(list(x), fps=args.fps).fl_image(make_uint8)
tt.write_videofile(os.path.join(args.save_path, 'input', fn), logger=None)
tt.close()
d += 1
postfix_kwargs = {}
if args.selfsup_loss:
postfix_kwargs = {'SelfsupLoss': selfsup_losses.avg, 'SupLoss': sup_losses.avg}
if args.selfsup_loss == 'sort' or args.selfsup_loss == 'fps':
postfix_kwargs['SelfsupAcc'] = selfsup_acc.avg
t.set_postfix(
DataTime=data_time.avg,
BatchTime=batch_time.avg,
Loss=losses.avg,
Acc=accs.avg,
**postfix_kwargs
)
if args.get_features:
torch.save(featsarr, os.path.join(args.save_path, 'input', 'features_and_fns.pt'))
if args.save_preds:
torch.save(predsarr, os.path.join(args.save_path, 'input', 'preds_and_fns.pt'))
if save_flag == True:
with open(os.path.join(args.save_path, 'results_{0:06}_{1:03}.csv'.format(args.local_rank, epoch)), 'w') as f:
wr = csv.writer(f)
wr.writerows(results)
if args.selfsup_loss == 'ctc':
return selfsup_losses.avg * -1, selfsup_losses.count
if accs.count > 0:
return accs.avg, accs.count
else:
return selfsup_acc.avg, selfsup_acc.count