train.py

import os

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

from utils import *
from metrics import *
import pickle
import argparse
import torch.optim.lr_scheduler as lr_scheduler
from model import SocialImplicit
from trajectory_augmenter import TrajectoryAugmenter
from CFG import CFG

parser = argparse.ArgumentParser()

#Social-Loss specific parameters
parser.add_argument('--w_norm',
                    type=float,
                    default=0.0001,
                    help='Intra-distance loss weight')
parser.add_argument('--w_cos',
                    type=float,
                    default=0.0001,
                    help='Angle between nodes loss weight')
parser.add_argument('--w_trip',
                    type=float,
                    default=0.0001,
                    help='Triplet loss weight')

#Data specifc paremeters
parser.add_argument('--obs_seq_len', type=int, default=8)
parser.add_argument('--pred_seq_len', type=int, default=12)
parser.add_argument('--dataset',
                    default='hotel',
                    help='eth,hotel,univ,zara1,zara2,sdd')

#Training specifc parameters
parser.add_argument('--batch_size',
                    type=int,
                    default=128,
                    help='minibatch size')
parser.add_argument('--num_epochs',
                    type=int,
                    default=50,
                    help='number of epochs')
parser.add_argument('--clip_grad',
                    type=float,
                    default=None,
                    help='gadient clipping')
parser.add_argument('--lr', type=float, default=0.01, help='learning rate')
parser.add_argument('--lr_sh_rate',
                    type=int,
                    default=45,
                    help='number of steps to drop the lr')

parser.add_argument('--tag', default='tag', help='personal tag for the model ')
args = parser.parse_args()

print('*' * 30)
print("Training initiating....")
print(args)

#Social-Loss
loss_store = {"l2": 0, "gl2": 0, "gcos": 0, "trip": 0}

_l1_mean = nn.L1Loss()


def cdist_cosine_sim(a, b, eps=1e-08):
    a_norm = a / torch.clamp(a.norm(dim=1)[:, None], min=eps)
    b_norm = b / torch.clamp(b.norm(dim=1)[:, None], min=eps)
    return torch.acos(
        torch.clamp(torch.mm(a_norm, b_norm.transpose(0, 1)),
                    min=-1.0 + eps,
                    max=1.0 - eps))


def reset_loss_store():
    global loss_store
    loss_store = {"l2": 0, "gl2": 0, "gcos": 0, "trip": 0}


def implicit_likelihood_estimation_fast_with_trip_geo(V_pred, V_target):

    V_pred = V_pred.contiguous()

    diff = torch.abs(V_pred - V_target)

    diff_sum = torch.sum(diff, dim=(1, 2, 3))
    _, indices = torch.sort(diff_sum)
    min_indx = indices[0]
    V_pred_min = V_pred[min_indx]
    V_target = V_target.squeeze()

    error = _l1_mean(V_pred_min, V_target)
    trip_loss = _l1_mean(V_pred_min, V_pred[indices[1]]) - _l1_mean(
        V_pred_min, V_pred[indices[-1]])

    V_pred_min_ = V_pred_min.reshape(-1, 2)
    V_target_ = V_target.reshape(-1, 2)

    #Geometric distance length
    norm_loss = torch.abs(
        torch.cdist(V_pred_min_.unsqueeze(0), V_pred_min_.unsqueeze(0), p=2.0)
        - torch.cdist(V_target_.unsqueeze(0), V_target_.unsqueeze(0), p=2.0)
    ).mean()

    #Gemometric distance angle
    cos_loss = torch.abs(
        cdist_cosine_sim(V_pred_min_, V_pred_min_) -
        cdist_cosine_sim(V_target_, V_target_)).mean()

    loss_store["l2"] += error.item()
    loss_store["gl2"] += norm_loss.item()
    loss_store["gcos"] += cos_loss.item()
    loss_store["trip"] += trip_loss.item()

    return error + args.w_norm * norm_loss + args.w_trip * trip_loss + args.w_cos * cos_loss


def graph_loss(V_pred, V_target, V_obs):
    return implicit_likelihood_estimation_fast_with_trip_geo(V_pred, V_target)


#Data prep
obs_seq_len = args.obs_seq_len
pred_seq_len = args.pred_seq_len
data_set = './datasets/' + args.dataset + '/'

dset_train = TrajectoryDataset(data_set + 'train/',
                               obs_len=obs_seq_len,
                               pred_len=pred_seq_len,
                               skip=1,
                               norm_lap_matr=True)

loader_train = DataLoader(
    dset_train,
    batch_size=1,  #This is irrelative to the args batch size parameter
    shuffle=True,
    num_workers=0)

dset_val = TrajectoryDataset(data_set + 'val/',
                             obs_len=obs_seq_len,
                             pred_len=pred_seq_len,
                             skip=1,
                             norm_lap_matr=True)

loader_val = DataLoader(
    dset_val,
    batch_size=1,  #This is irrelative to the args batch size parameter
    shuffle=False,
    num_workers=1)

#Defining the model
is_eth = args.dataset == 'eth'
if is_eth:
    noise_weight = CFG["noise_weight_eth"]
else:
    noise_weight = CFG["noise_weight"]

model = SocialImplicit(spatial_input=CFG["spatial_input"],
                       spatial_output=CFG["spatial_output"],
                       temporal_input=CFG["temporal_input"],
                       temporal_output=CFG["temporal_output"],
                       bins=CFG["bins"],
                       noise_weight=noise_weight).cuda().double()

#Optimizer and Schedule
optimizer = optim.SGD(model.parameters(), lr=args.lr)
scheduler = optim.lr_scheduler.StepLR(optimizer,
                                      step_size=args.lr_sh_rate,
                                      gamma=0.1)

#Check pointing
checkpoint_dir = './checkpoint/' + args.tag + '/'

if not os.path.exists(checkpoint_dir):
    os.makedirs(checkpoint_dir)

with open(checkpoint_dir + 'args.pkl', 'wb') as fp:
    pickle.dump(args, fp)

print('Data and model loaded')
print('Checkpoint dir:', checkpoint_dir)

#Training
metrics = {'train_loss': [], 'val_loss': []}
constant_metrics = {'min_val_epoch': -1, 'min_val_loss': 9999999999999999}
trajaugmenter = TrajectoryAugmenter(data_loader=loader_train)


def train(epoch):
    global metrics, loader_train, loss_store
    model.train()

    total_loss = 0
    batch_loss = 0
    for cnt, batch in enumerate(loader_train):

        #Get data
        obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_gt_rel, non_linear_ped, loss_mask, V_obs, A_obs, V_tr, A_tr = batch

        V_obs, V_tr, obs_traj, pred_traj_gt = trajaugmenter.augment(
            V_obs, V_tr, obs_traj, pred_traj_gt)

        V_obs, V_tr, A_obs, obs_traj = V_obs.cuda().double(), V_tr.cuda(
        ).double(), A_obs.cuda().double(), obs_traj.cuda().double()

        optimizer.zero_grad()

        #Forward
        V_pred = model(V_obs.permute(0, 3, 1, 2), obs_traj)
        V_pred = V_pred.permute(0, 2, 3, 1)

        #Loss
        batch_loss += graph_loss(V_pred, V_tr, V_obs)
        total_loss += batch_loss.item()

        #Learn
        if cnt % args.batch_size == 0 and cnt != 0:
            batch_loss = batch_loss / args.batch_size
            batch_loss.backward()
            if args.clip_grad is not None:
                torch.nn.utils.clip_grad_norm_(model.parameters(),
                                               args.clip_grad)
            optimizer.step()
            #Log
            print(args.tag, ' |TRAIN:', '\t Epoch:', epoch, '\t Batch loss:',
                  batch_loss.item())

            loss_store["l2"] /= args.batch_size
            loss_store["gl2"] /= args.batch_size
            loss_store["gcos"] /= args.batch_size
            print("Detailed train loss:", loss_store)
            reset_loss_store()

            #Reset
            batch_loss = 0
    metrics['train_loss'].append(total_loss / (cnt + 1))


iteration = 0


def vald():
    global metrics, loader_val, constant_metrics, iteration, loss_store
    model.eval()
    total_loss = 0

    with torch.no_grad():
        for cnt, batch in enumerate(loader_val):

            #Get data
            obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_gt_rel, non_linear_ped, loss_mask, V_obs, A_obs, V_tr, A_tr = batch

            #Forward
            V_obs, V_tr, A_obs, obs_traj = V_obs.cuda().double(), V_tr.cuda(
            ).double(), A_obs.cuda().double(), obs_traj.cuda().double()

            V_pred = model(V_obs.permute(0, 3, 1, 2), obs_traj)
            V_pred = V_pred.permute(0, 2, 3, 1)

            #Loss
            total_loss += graph_loss(V_pred, V_tr, V_obs).item()

        print(args.tag, ' |VALD:', '\t Iteration:', iteration, '\t Loss:',
              total_loss / (cnt + 1))
        metrics['val_loss'].append(total_loss / (cnt + 1))
        loss_store["l2"] /= (cnt + 1)
        loss_store["gl2"] /= (cnt + 1)
        loss_store["gcos"] /= (cnt + 1)
        print("Detailed val loss:", loss_store)
        reset_loss_store()
        store_per = 0.05 * constant_metrics['min_val_loss']

        if (constant_metrics['min_val_loss'] -
                metrics['val_loss'][-1]) > store_per:
            # if metrics['val_loss'][-1] < constant_metrics['min_val_loss']:
            constant_metrics['min_val_loss'] = metrics['val_loss'][-1]
            constant_metrics['min_val_epoch'] = iteration
            torch.save(model.state_dict(),
                       checkpoint_dir + 'val_best.pth')  # OK
    iteration += 1


print('Training started ...')
for epoch in range(args.num_epochs):
    train(epoch)
    vald()

    scheduler.step()

    print('*' * 30)
    print(args.tag, ' |Epoch:', args.tag, ":", epoch)
    for k, v in metrics.items():
        if len(v) > 0:
            print(k, v[-1])

    print(constant_metrics)
    print('*' * 30)
    for g in optimizer.param_groups:
        print("***------------->LR = ", g['lr'])
    with open(checkpoint_dir + 'metrics.pkl', 'wb') as fp:
        pickle.dump(metrics, fp)

    with open(checkpoint_dir + 'constant_metrics.pkl', 'wb') as fp:
        pickle.dump(constant_metrics, fp)