96 lines
4.2 KiB
Python
96 lines
4.2 KiB
Python
|
import torch
|
||
|
|
import wandb
|
||
|
|
from tqdm import tqdm
|
||
|
|
from trainer.train_builder import TRAINER
|
||
|
|
|
||
|
|
|
||
|
|
@TRAINER.register("OAD")
|
||
|
|
def train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None):
|
||
|
|
epoch_loss = 0
|
||
|
|
for it, (vid_name, rgb_input, flow_input, target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')):
|
||
|
|
rgb_input, flow_input, target = rgb_input.cuda(), flow_input.cuda(), target.cuda()
|
||
|
|
model.train()
|
||
|
|
if scaler != None:
|
||
|
|
with torch.cuda.amp.autocast():
|
||
|
|
out_dict = model(rgb_input, flow_input)
|
||
|
|
loss = criterion(out_dict, target)
|
||
|
|
optimizer.zero_grad(set_to_none=True)
|
||
|
|
scaler.scale(loss).backward()
|
||
|
|
scaler.step(optimizer)
|
||
|
|
scaler.update()
|
||
|
|
# scaler is None as argparse by default
|
||
|
|
else:
|
||
|
|
out_dict = model(rgb_input, flow_input)
|
||
|
|
loss = criterion(out_dict, target)
|
||
|
|
optimizer.zero_grad(set_to_none=True)
|
||
|
|
loss.backward()
|
||
|
|
optimizer.step()
|
||
|
|
|
||
|
|
epoch_loss += loss.item()
|
||
|
|
if use_wandb == True:
|
||
|
|
wandb.log({"Train Loss": loss.item()})
|
||
|
|
|
||
|
|
if writer != None:
|
||
|
|
writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader))
|
||
|
|
return epoch_loss
|
||
|
|
|
||
|
|
@TRAINER.register("OADTR")
|
||
|
|
def train_one_epoch_oadtr(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None):
|
||
|
|
epoch_loss = 0
|
||
|
|
for it, (vid_name, rgb_input, flow_input, enc_target, distance_target, class_h_target, dec_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')):
|
||
|
|
rgb_input, flow_input, enc_target, distance_target, class_h_target = rgb_input.cuda(), flow_input.cuda(), enc_target.cuda(), distance_target.cuda(), class_h_target.cuda()
|
||
|
|
model.train()
|
||
|
|
if scaler != None:
|
||
|
|
with torch.cuda.amp.autocast():
|
||
|
|
out_dict = model(rgb_input, flow_input)
|
||
|
|
loss = criterion(out_dict, enc_target, dec_target)
|
||
|
|
optimizer.zero_grad(set_to_none=True)
|
||
|
|
scaler.scale(loss).backward()
|
||
|
|
scaler.step(optimizer)
|
||
|
|
scaler.update()
|
||
|
|
else:
|
||
|
|
out_dict = model(rgb_input, flow_input)
|
||
|
|
loss = criterion(out_dict, enc_target, dec_target)
|
||
|
|
optimizer.zero_grad(set_to_none=True)
|
||
|
|
loss.backward()
|
||
|
|
optimizer.step()
|
||
|
|
|
||
|
|
epoch_loss += loss.item()
|
||
|
|
if use_wandb == True:
|
||
|
|
wandb.log({"Train Loss": loss.item()})
|
||
|
|
|
||
|
|
if writer != None:
|
||
|
|
writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader))
|
||
|
|
|
||
|
|
return epoch_loss
|
||
|
|
|
||
|
|
|
||
|
|
@TRAINER.register("ANTICIPATION")
|
||
|
|
def ant_train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None):
|
||
|
|
epoch_loss = 0
|
||
|
|
for it, (vid_name, rgb_input, flow_input, target, ant_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')):
|
||
|
|
rgb_input, flow_input, target, ant_target = rgb_input.cuda(), flow_input.cuda(), target.cuda(), ant_target.cuda()
|
||
|
|
model.train()
|
||
|
|
if scaler != None:
|
||
|
|
with torch.cuda.amp.autocast():
|
||
|
|
out_dict = model(rgb_input, flow_input)
|
||
|
|
loss = criterion(out_dict, target, ant_target)
|
||
|
|
optimizer.zero_grad(set_to_none=True)
|
||
|
|
scaler.scale(loss).backward()
|
||
|
|
scaler.step(optimizer)
|
||
|
|
scaler.update()
|
||
|
|
else:
|
||
|
|
out_dict = model(rgb_input, flow_input)
|
||
|
|
loss = criterion(out_dict, target, ant_target)
|
||
|
|
optimizer.zero_grad(set_to_none=True)
|
||
|
|
loss.backward()
|
||
|
|
optimizer.step()
|
||
|
|
epoch_loss += loss.item()
|
||
|
|
|
||
|
|
if use_wandb == True:
|
||
|
|
wandb.log({"Train Loss": loss.item()})
|
||
|
|
|
||
|
|
if writer != None:
|
||
|
|
writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader))
|
||
|
|
|
||
|
|
return epoch_loss
|