train_abs.py

import os
import sys
import random
import time
import cv2
import argparse
import numpy as np
import matplotlib.pyplot as plt

import torch.multiprocessing as mp

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
from torch.optim.lr_scheduler import ReduceLROnPlateau

from models_abs import PropagationNetwork
from data_abs import PhysicsCLEVRDataset, collate_fn

from utils import count_parameters, Tee
import pdb
from utils_tube import set_debugger 
set_debugger()

parser = argparse.ArgumentParser()
parser.add_argument('--pn', type=int, default=1)
parser.add_argument('--pstep', type=int, default=2)

parser.add_argument('--data_dir', default='')
parser.add_argument('--label_dir', default='')

parser.add_argument('--n_rollout', type=int, default=0)
parser.add_argument('--n_particle', type=int, default=0)
parser.add_argument('--time_step', type=int, default=0)
parser.add_argument('--nf_relation', type=int, default=128)
parser.add_argument('--nf_particle', type=int, default=128)
parser.add_argument('--nf_effect', type=int, default=128*4)
parser.add_argument('--env', default='CLEVR')
parser.add_argument('--dt', default=1./5.)
parser.add_argument('--train_valid_ratio', type=float, default=0.90909)
parser.add_argument('--outf', default='files')
parser.add_argument('--dataf', default='data')
parser.add_argument('--num_workers', type=int, default=4)
parser.add_argument('--log_per_iter', type=int, default=1000)
parser.add_argument('--ckp_per_iter', type=int, default=50000)
parser.add_argument('--eval', type=int, default=0)
parser.add_argument('--edge_superv', type=int, default=1, help='whether to include edge supervision')
parser.add_argument('--use_attr', type=int, default=1, help='whether using attributes or not')

parser.add_argument('--n_his', type=int, default=2)
parser.add_argument('--frame_offset', type=int, default=5)
parser.add_argument('--gen_valid_idx', type=int, default=0)

parser.add_argument('--lam_mask', type=float, default=0.2)
parser.add_argument('--lam_position', type=float, default=20.)
parser.add_argument('--lam_image', type=float, default=0.25)
parser.add_argument('--lam_collision', type=float, default=0.6)
parser.add_argument('--lam_hw', type=float, default=20.)

parser.add_argument('--n_epoch', type=int, default=1000)
parser.add_argument('--beta1', type=float, default=0.9)
parser.add_argument('--lr', type=float, default=0.0001)
parser.add_argument('--batch_size', type=int, default=1)
parser.add_argument('--forward_times', type=int, default=2)

parser.add_argument('--resume_epoch', type=int, default=0)
parser.add_argument('--resume_iter', type=int, default=0)

# object attributes (material, shape):
# [rubber, metal, cube, cylinder, sphere]
parser.add_argument('--attr_dim', type=int, default=5)

# object state:
# [mask, dx, dy, dw, dh, r, g, b]
parser.add_argument('--state_dim', type=int, default=6)

# relation:
# [collision, dx, dy]
parser.add_argument('--relation_dim', type=int, default=3)
parser.add_argument('--tube_mode', type=int, default=0)

# new input
parser.add_argument('--debug', type=int, default=0)
parser.add_argument('--box_only_flag', type=int, default=0)
parser.add_argument('--add_hw_state_flag', type=int, default=0)
parser.add_argument('--rm_mask_state_flag', type=int, default=0)
parser.add_argument('--add_xyhw_state_flag', type=int, default=0)

args = parser.parse_args()

### deal with issue that dataloader hangs
cv2.setNumThreads(0)

if args.env == 'CLEVR':
    #args.n_rollout = 11000
    #args.time_step = 128
    #args.n_rollout = 11000
    #args.train_valid_ratio = 0.909
    if args.debug:
        args.n_rollout = 15
        args.train_valid_ratio = 0.666
        shuffle_flag = False
    else:
        args.n_rollout = 15000
        args.train_valid_ratio = 0.666
        shuffle_flag = True
else:
    raise AssertionError("Unsupported env")

args.outf = args.outf + '_' + args.env
if args.use_attr == 0:
    args.outf += '_noAttr'
if args.edge_superv == 0:
    args.outf += '_noEdgeSuperv'
if args.pn:
    args.outf += '_pn'
args.outf += '_pstep_' + str(args.pstep) + '_abs'
# args.dataf = args.dataf + '_' + args.env


os.system('mkdir -p ' + args.outf)
# os.system('mkdir -p ' + args.dataf)

# setup recorder
#tee = Tee(os.path.join(args.outf, 'train.log'), 'w')
print(args)

# generate data
datasets = {phase: PhysicsCLEVRDataset(args, phase)  for phase in ['train', 'valid']}
use_gpu = torch.cuda.is_available()

shuffle_flag=True
num_workers = args.num_workers
if args.debug:
    shuffle_flag=False 
    num_workers = 0

dataloaders = {x: torch.utils.data.DataLoader(
    datasets[x], batch_size=args.batch_size,
    shuffle=shuffle_flag if x == 'train' else False,
    num_workers=num_workers,
    collate_fn=collate_fn)
    for x in ['train', 'valid']}

# define propagation network
model = PropagationNetwork(args, residual=True, use_gpu=use_gpu)

print("model #params: %d" % count_parameters(model))

if args.resume_epoch > 0 or args.resume_iter > 0:
    model_path = os.path.join(args.outf, 'net_epoch_%d_iter_%d.pth' % (args.resume_epoch, args.resume_iter))
    print("Loading saved ckp from %s" % model_path)
    model.load_state_dict(torch.load(model_path))

# criterion
criterionMSE = nn.MSELoss()

# optimizer
optimizer = optim.Adam(model.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
scheduler = ReduceLROnPlateau(optimizer, 'min', factor=0.8, patience=3, verbose=True)

if use_gpu:
    model = model.cuda()
    criterionMSE = criterionMSE.cuda()


st_epoch = args.resume_epoch if args.resume_epoch > 0 else 0
best_valid_loss = np.inf

for epoch in range(st_epoch, args.n_epoch):

    phases = ['train', 'valid'] if args.eval == 0 else ['valid']
    for phase in phases:

        model.train(phase=='train')

        losses = 0.
        losses_mask = 0.
        losses_position = 0.
        losses_image = 0.
        losses_collision = 0.
        losses_hw = 0.
        for i, data in enumerate(dataloaders[phase]):
            attr, x, rel, label_obj, label_rel = data

            node_r_idx, node_s_idx, Ra = rel[3], rel[4], rel[5]
            Rr_idx, Rs_idx, value = rel[0], rel[1], rel[2]

            Rr = torch.sparse.FloatTensor(
                Rr_idx, value, torch.Size([node_r_idx.shape[0], value.size(0)]))
            Rs = torch.sparse.FloatTensor(
                Rs_idx, value, torch.Size([node_s_idx.shape[0], value.size(0)]))

            data = [attr, x, Rr, Rs, Ra, label_obj, label_rel]
            with torch.set_grad_enabled(phase=='train'):
                for d in range(len(data)):
                    if use_gpu:
                        data[d] = Variable(data[d].cuda())
                    else:
                        data[d] = Variable(data[d])

                attr, x, Rr, Rs, Ra, label_obj, label_rel = data
                pred_obj, pred_rel = model(
                    attr, x, Rr, Rs, Ra, node_r_idx, node_s_idx, args.pstep)

            if args.box_only_flag:
                position = pred_obj
                collision = pred_rel
            else:
                mask = pred_obj[:, 0]
                position = pred_obj[:, 1:3]
                image = pred_obj[:, 3:6]
                collision = pred_rel
            if args.add_hw_state_flag or args.add_xyhw_state_flag:
                hw_pred = pred_obj[:, 6:]

            '''
            print('mask\n', mask)
            print('x\n', position[0])
            print('y\n', position[1])
            print('img\n', image[0])
            '''
            loss = 0.0
            if args.box_only_flag:
                loss_position = criterionMSE(position, label_obj)
                loss_collision = criterionMSE(collision, label_rel)
                
                loss_image = torch.zeros(1).cuda()
                loss_mask = torch.zeros(1).cuda()
                losses_image += np.sqrt(loss_image.item())
                losses_mask += np.sqrt(loss_mask.item())
            else:
                loss_position = criterionMSE(position, label_obj[:, 1:3])
                if args.rm_mask_state_flag:
                    loss_mask = torch.zeros(1).cuda()
                else:
                    loss_mask = criterionMSE(mask, label_obj[:, 0])
                loss_image = criterionMSE(image, label_obj[:, 3:6])
                loss_collision = criterionMSE(collision, label_rel)
            
                losses_mask += np.sqrt(loss_mask.item())
                losses_image += np.sqrt(loss_image.item())
                loss = loss_mask * args.lam_mask
                loss += loss_image * args.lam_image
           
            if args.add_hw_state_flag or args.add_xyhw_state_flag:
                loss_hw = criterionMSE(hw_pred, label_obj[:, 6:])
                loss += loss_hw * args.lam_hw
                losses_hw += np.sqrt(loss_hw.item())

            losses_position += np.sqrt(loss_position.item())
            losses_collision += np.sqrt(loss_collision.item())
            loss += loss_position * args.lam_position

            if args.edge_superv:
                loss += loss_collision * args.lam_collision

            losses += np.sqrt(loss.item())
            #pdb.set_trace()
            if phase == 'train':
                if i % args.forward_times == 0:
                    if i != 0:
                        loss_acc /= args.forward_times
                        optimizer.zero_grad()
                        loss_acc.backward()
                        optimizer.step()
                    loss_acc = loss
                else:
                    loss_acc += loss

            if i % args.log_per_iter == 0:
                log = '%s [%d/%d][%d/%d] Loss: %.6f %.6f %.6f %.6f %.6f, Agg: %.6f %.6f %.6f %.6f %.6f' % \
                      (phase, epoch, args.n_epoch, i, len(dataloaders[phase]),
                       np.sqrt(loss_mask.item()), np.sqrt(loss_position.item()),
                       np.sqrt(loss_image.item()), np.sqrt(loss_collision.item()),
                       np.sqrt(loss.item()),
                       losses_mask / (i + 1), losses_position / (i + 1),
                       losses_image / (i + 1), losses_collision / (i + 1),
                       losses / (i + 1))

                print(log)
                if args.add_hw_state_flag:
                    print('hw loss: %.6f, agg: %.6f' %(np.sqrt(loss_hw.item()), losses_hw/ (i+1)))
            if phase == 'train' and i > 0 and i % args.ckp_per_iter == 0:
                torch.save(model.state_dict(), '%s/net_epoch_%d_iter_%d.pth' % (args.outf, epoch, i))
        losses /= len(dataloaders[phase])
        log = '%s [%d/%d] Loss: %.4f, Best valid: %.4f' % \
              (phase, epoch, args.n_epoch, losses, best_valid_loss)
        print(log)

        if phase == 'valid':
            scheduler.step(losses)
            if(losses < best_valid_loss):
                best_valid_loss = losses
                torch.save(model.state_dict(), '%s/net_best.pth' % (args.outf))