main_auxi_v0_noBAN.py

#!/usr/bin/env python
# coding=utf-8
'''
Author:Tai Lei
Date:Wed Sep 19 20:30:48 2018
Info:
References: https://github.com/pytorch/examples/tree/master/imagenet
'''
import argparse
import os
import random
import time
import datetime
import math
import logging

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim as optim
from tensorboardX import SummaryWriter

import numpy as np

import sys
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))
import auxi_ver0.auxi_net_v0 as resnet_carla
from auxi_ver0.carla_loader_db_auxi_v0 import CarlaH5Data
from auxi_ver0.helper_auxi_v0 import AverageMeter, save_checkpoint

parser = argparse.ArgumentParser(description='Carla CIL training')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--batch_size', default=200, type=int, metavar='N',
                    help='batch size of training')
parser.add_argument('--speed_weight', default=1, type=float,
                    help='speed weight')
parser.add_argument('--branch-weight', default=1, type=float,
                    help='branch weight')
parser.add_argument('--id', default="training", type=str)
parser.add_argument('--train-dir', default="/SSD1/datasets/carla/additional_db/val/",
                    type=str, metavar='PATH', help='training dataset')
parser.add_argument('--eval-dir', default="/SSD1/datasets/carla/additional_db/val/",
                    type=str, metavar='PATH',
                    help='evaluation dataset')
parser.add_argument('--epochs', default=100, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('--lr', '--learning-rate', default=1e-4, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq', '-p', default=10, type=int,
                    metavar='N', help='print frequency (default: 10)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                    help='evaluate model on validation set')
parser.add_argument('--evaluate-log', default="",
                    type=str, metavar='PATH',
                    help='path to log evaluation results (default: none)')
parser.add_argument('--world-size', default=1, type=int,
                    help='number of distributed processes')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=2, type=int,
                    help='GPU id to use.')



def output_log(output_str, logger=None):
    """
    standard output and logging
    """
    print("[{}]: {}".format(datetime.datetime.now(), output_str))
    if logger is not None:
        logger.critical("[{}]: {}".format(datetime.datetime.now(), output_str))


def log_args(logger):
    '''
    log args
    '''
    attrs = [(p, getattr(args, p)) for p in dir(args) if not p.startswith('_')]
    for key, value in attrs:
        output_log("{}: {}".format(key, value), logger=logger)


def main():
    global args
    args = parser.parse_args()
    log_dir = os.path.join("./", "logs", args.id)
    run_dir = os.path.join("./", "runs", args.id)
    save_weight_dir = os.path.join("./save_models", args.id)
    os.makedirs(log_dir, exist_ok=True)
    os.makedirs(save_weight_dir, exist_ok=True)

    logging.basicConfig(filename=os.path.join(log_dir, "carla_training.log"),
                        level=logging.ERROR)
    tsbd = SummaryWriter(log_dir=run_dir)
    log_args(logging)
    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        output_log(
            'You have chosen to seed training. '
            'This will turn on the CUDNN deterministic setting, '
            'which can slow down your training considerably! '
            'You may see unexpected behavior when restarting '
            'from checkpoints.', logger=logging)

    if args.gpu is not None:
        output_log('You have chosen a specific GPU. This will completely '
                   'disable data parallelism.', logger=logging)

    args.distributed = args.world_size > 1

    if args.distributed:
        dist.init_process_group(backend=args.dist_backend,
                                init_method=args.dist_url,
                                world_size=args.world_size,
                                rank=0)

    model = resnet_carla.resnet34_carla(True)
    # criterion = EgoLoss()
    criterion = nn.MSELoss()

    tsbd.add_graph(model,
                   (torch.zeros(1, 3, 88, 200),
                    torch.zeros(1, 1)))

    if args.gpu is not None:
        model = model.cuda(args.gpu)
    else:
        model = torch.nn.DataParallel(model).cuda()

    # TODO check other papers optimizers
    optimizer = optim.Adam(
        model.parameters(), args.lr, betas=(0.7, 0.85))
    lr_scheduler = optim.lr_scheduler.StepLR(
        optimizer, step_size=7, gamma=0.5)

    # optionally resume from a checkpoint
    if args.resume:
        args.resume = os.path.join(save_weight_dir, args.resume)
        if os.path.isfile(args.resume):
            output_log("=> loading checkpoint '{}'".format(args.resume),
                       logging)
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            lr_scheduler.load_state_dict(checkpoint['scheduler'])
            output_log("=> loaded checkpoint '{}' (epoch {})"
                       .format(args.resume, checkpoint['epoch']), logging)
        else:
            output_log("=> no checkpoint found at '{}'".format(args.resume),
                       logging)

    cudnn.benchmark = True

    carla_data = CarlaH5Data(
        train_folder=args.train_dir,
        eval_folder=args.eval_dir,
        batch_size=args.batch_size,
        num_workers=args.workers)

    train_loader = carla_data.loaders["train"]
    eval_loader = carla_data.loaders["eval"]
    best_prec = math.inf

    if args.evaluate:
        args.id = args.id+"_test"
        if not os.path.isfile(args.resume):
            output_log("=> no checkpoint found at '{}'"
                       .format(args.resume), logging)
            return
        if args.evaluate_log == "":
            output_log("=> please set evaluate log path with --evaluate-log <log-path>")

        # TODO add test func
        evaluate(eval_loader, model, criterion, 0, tsbd)
        return

    for epoch in range(args.start_epoch, args.epochs):
        branch_losses, speed_losses, losses = \
            train(train_loader, model, criterion, optimizer, epoch, tsbd)

        #prec = evaluate(eval_loader, model, criterion, epoch, tsbd)
        prec = 0

        lr_scheduler.step()

        # remember best prec@1 and save checkpoint
        is_best = prec < best_prec
        best_prec = min(prec, best_prec)
        save_checkpoint(
            {'epoch': epoch + 1,
             'state_dict': model.state_dict(),
             'best_prec': best_prec,
             'scheduler': lr_scheduler.state_dict(),
             'optimizer': optimizer.state_dict()},
            args.id,
            is_best,
            os.path.join(
                save_weight_dir,
                "{}_{}.pth".format(epoch+1, args.id))
            )

def acos_safe(x, eps=1e-4):
    sign = torch.sign(x)
    slope = np.arccos(1-eps) / eps
    return torch.where(abs(x) <= 1-eps,
                    torch.acos(x),
                    torch.acos(sign * (1 - eps)) - slope*sign*(abs(x) - 1 + eps))

def get_predicted_wheel_location(steering_angle, v, time_stamp=0.1):
    wheel_heading = steering_angle
    wheel_traveled_dis = v * time_stamp

    return [wheel_traveled_dis * torch.cos(wheel_heading), wheel_traveled_dis * torch.sin(wheel_heading)]

def get_predicted_steering(pred_x, pred_y, v, time_stamp=0.1):
    # to prevent over range of arccos or arcsin
    # eps = 1e-7
    eps = 0.00001
    cal_x = torch.clamp(pred_x / (v * time_stamp), -1 + eps, 1 - eps)
    cal_y = torch.clamp(pred_y / (v * time_stamp), -1 + eps, 1 - eps)

    steering_angle_x = torch.acos(cal_x)
    steering_angle_y = torch.asin(cal_y)

    return [steering_angle_x, steering_angle_y]

def get_predicted_velocity(pred_x, pred_y, steering_angle, time_stamp=0.1):
    v_x = pred_x / (time_stamp * torch.cos(steering_angle))
    v_y = pred_y / (time_stamp * torch.sin(steering_angle))

    return [v_x, v_y]

def train(loader, model, criterion, optimizer, epoch, writer):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    branch_losses = AverageMeter()
    speed_losses = AverageMeter()

    # switch to train mode
    model.train()
    end = time.time()
    step = epoch * len(loader)
    for i, (img, speed, target, mask) in enumerate(loader):
        data_time.update(time.time() - end)

        # if args.gpu is not None:
        img = img.cuda(args.gpu, non_blocking=True)
        speed = speed.cuda(args.gpu, non_blocking=True)
        target = target.cuda(args.gpu, non_blocking=True)
        mask = mask.cuda(args.gpu, non_blocking=True)

        branches_out, pred_speed = model(img, speed)

        mask_out = branches_out * mask
        branch_loss = criterion(mask_out, target) * 4
        speed_loss = criterion(pred_speed, speed)

        loss = args.branch_weight * branch_loss + \
            args.speed_weight * speed_loss

        losses.update(loss.item(), args.batch_size)
        branch_losses.update(branch_loss.item(), args.batch_size)
        speed_losses.update(speed_loss.item(), args.batch_size)

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0 or i == len(loader):
            writer.add_scalar('train/branch_loss', branch_losses.val, step+i)
            writer.add_scalar('train/speed_loss', speed_losses.val, step+i)
            writer.add_scalar('train/loss', losses.val, step+i)
            output_log(
                'Epoch: [{0}][{1}/{2}]\t'
                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                'Branch loss {branch_loss.val:.3f} ({branch_loss.avg:.3f})\t'
                'Speed loss {speed_loss.val:.3f} ({speed_loss.avg:.3f})\t'
                'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                .format(
                    epoch, i, len(loader), batch_time=batch_time,
                    data_time=data_time, branch_loss=branch_losses,
                    speed_loss=speed_losses, loss=losses), logging)

    return branch_losses.avg, speed_losses.avg, losses.avg


def evaluate(loader, model, criterion, epoch, writer):
    batch_time = AverageMeter()
    losses = AverageMeter()

    # switch to evaluate mode
    model.eval()
    step = epoch * len(loader)
    with torch.no_grad():
        end = time.time()
        for i, (img, speed, target_lateral, mask_lateral, target_longi, mask_longi, command) in enumerate(loader):
            img = img.cuda(args.gpu, non_blocking=True)
            speed = speed.cuda(args.gpu, non_blocking=True)
            target_lateral = target_lateral.cuda(args.gpu, non_blocking=True)
            mask_lateral = mask_lateral.cuda(args.gpu, non_blocking=True)
            target_longi = target_longi.cuda(args.gpu, non_blocking=True)
            mask_longi = mask_longi.cuda(args.gpu, non_blocking=True)
            command = command.cuda(args.gpu, non_blocking=True)

            pred_speed, action_longi, action_lateral = model(img, speed)
            speed_loss = criterion(pred_speed, speed)
            mask_out_longi = action_longi * mask_longi
            branch_longi_loss = criterion(mask_out_longi, target_longi) * 4

            mask_out_lateral = action_lateral * mask_lateral
            branch_lateral_loss = criterion(mask_out_lateral, target_lateral) * 4

            loss = args.branch_weight * branch_longi_loss + \
                   args.branch_weight * branch_lateral_loss + \
                   args.speed_weight * speed_loss

            # measure accuracy and record loss
            losses.update(loss.item(), args.batch_size)

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % args.print_freq == 0 or i == len(loader):
                writer.add_scalar('eval/loss', losses.val, step+i)
                output_log(
                  'Test: [{0}/{1}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  .format(
                      i, len(loader), batch_time=batch_time,
                      loss=losses), logging)
    return losses.avg


if __name__ == '__main__':
    main()