def train(conf, train_shape_list, train_data_list, val_data_list,
all_train_data_list):
# create training and validation datasets and data loaders
data_features = ['pcs', 'pc_pxids', 'pc_movables', 'gripper_img_target', 'gripper_direction_camera', 'gripper_forward_direction_camera', \
'result', 'cur_dir', 'shape_id', 'trial_id', 'is_original']
# load network model
model_def = utils.get_model_module(conf.model_version)
# create models
network = model_def.Network(conf.feat_dim)
utils.printout(conf.flog, '\n' + str(network) + '\n')
# create optimizers
network_opt = torch.optim.Adam(network.parameters(),
lr=conf.lr,
weight_decay=conf.weight_decay)
# learning rate scheduler
network_lr_scheduler = torch.optim.lr_scheduler.StepLR(
network_opt, step_size=conf.lr_decay_every, gamma=conf.lr_decay_by)
# create logs
if not conf.no_console_log:
header = ' Time Epoch Dataset Iteration Progress(%) LR TotalLoss'
if not conf.no_tb_log:
# https://github.com/lanpa/tensorboard-pytorch
from tensorboardX import SummaryWriter
train_writer = SummaryWriter(os.path.join(conf.exp_dir, 'train'))
val_writer = SummaryWriter(os.path.join(conf.exp_dir, 'val'))
# send parameters to device
network.to(conf.device)
utils.optimizer_to_device(network_opt, conf.device)
# load dataset
train_dataset = SAPIENVisionDataset([conf.primact_type], conf.category_types, data_features, conf.buffer_max_num, \
abs_thres=conf.abs_thres, rel_thres=conf.rel_thres, dp_thres=conf.dp_thres, img_size=conf.img_size, no_true_false_equal=conf.no_true_false_equal)
val_dataset = SAPIENVisionDataset([conf.primact_type], conf.category_types, data_features, conf.buffer_max_num, \
abs_thres=conf.abs_thres, rel_thres=conf.rel_thres, dp_thres=conf.dp_thres, img_size=conf.img_size, no_true_false_equal=conf.no_true_false_equal)
val_dataset.load_data(val_data_list)
utils.printout(conf.flog, str(val_dataset))
val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=conf.batch_size, shuffle=False, pin_memory=True, \
num_workers=0, drop_last=True, collate_fn=utils.collate_feats, worker_init_fn=utils.worker_init_fn)
val_num_batch = len(val_dataloader)
# create a data generator
datagen = DataGen(conf.num_processes_for_datagen, conf.flog)
# sample succ
if conf.sample_succ:
sample_succ_list = []
sample_succ_dirs = []
# start training
start_time = time.time()
last_train_console_log_step, last_val_console_log_step = None, None
# if resume
start_epoch = 0
if conf.resume:
# figure out the latest epoch to resume
for item in os.listdir(os.path.join(conf.exp_dir, 'ckpts')):
if item.endswith('-train_dataset.pth'):
start_epoch = int(item.split('-')[0])
# load states for network, optimizer, lr_scheduler, sample_succ_list
data_to_restore = torch.load(
os.path.join(conf.exp_dir, 'ckpts',
'%d-network.pth' % start_epoch))
network.load_state_dict(data_to_restore)
data_to_restore = torch.load(
os.path.join(conf.exp_dir, 'ckpts',
'%d-optimizer.pth' % start_epoch))
network_opt.load_state_dict(data_to_restore)
data_to_restore = torch.load(
os.path.join(conf.exp_dir, 'ckpts',
'%d-lr_scheduler.pth' % start_epoch))
network_lr_scheduler.load_state_dict(data_to_restore)
# rmdir and make a new dir for the current sample-succ directory
old_sample_succ_dir = os.path.join(
conf.data_dir, 'epoch-%04d_sample-succ' % (start_epoch - 1))
utils.force_mkdir(old_sample_succ_dir)
# train for every epoch
for epoch in range(start_epoch, conf.epochs):
### collect data for the current epoch
if epoch > start_epoch:
utils.printout(
conf.flog,
f' [{strftime("%H:%M:%S", time.gmtime(time.time()-start_time)):>9s} Waiting epoch-{epoch} data ]'
)
train_data_list = datagen.join_all()
utils.printout(
conf.flog,
f' [{strftime("%H:%M:%S", time.gmtime(time.time()-start_time)):>9s} Gathered epoch-{epoch} data ]'
)
cur_data_folders = []
for item in train_data_list:
item = '/'.join(item.split('/')[:-1])
if item not in cur_data_folders:
cur_data_folders.append(item)
for cur_data_folder in cur_data_folders:
with open(os.path.join(cur_data_folder, 'data_tuple_list.txt'),
'w') as fout:
for item in train_data_list:
if cur_data_folder == '/'.join(item.split('/')[:-1]):
fout.write(item.split('/')[-1] + '\n')
# load offline-generated sample-random data
for item in all_train_data_list:
valid_id_l = conf.num_interaction_data_offline + conf.num_interaction_data * (
epoch - 1)
valid_id_r = conf.num_interaction_data_offline + conf.num_interaction_data * epoch
if valid_id_l <= int(item.split('_')[-1]) < valid_id_r:
train_data_list.append(item)
### start generating data for the next epoch
# sample succ
if conf.sample_succ:
if conf.resume and epoch == start_epoch:
sample_succ_list = torch.load(
os.path.join(conf.exp_dir, 'ckpts',
'%d-sample_succ_list.pth' % start_epoch))
else:
torch.save(
sample_succ_list,
os.path.join(conf.exp_dir, 'ckpts',
'%d-sample_succ_list.pth' % epoch))
for item in sample_succ_list:
datagen.add_one_recollect_job(item[0], item[1], item[2],
item[3], item[4], item[5],
item[6])
sample_succ_list = []
sample_succ_dirs = []
cur_sample_succ_dir = os.path.join(
conf.data_dir, 'epoch-%04d_sample-succ' % epoch)
utils.force_mkdir(cur_sample_succ_dir)
# start all jobs
datagen.start_all()
utils.printout(
conf.flog,
f' [ {strftime("%H:%M:%S", time.gmtime(time.time()-start_time)):>9s} Started generating epoch-{epoch+1} data ]'
)
### load data for the current epoch
if conf.resume and epoch == start_epoch:
train_dataset = torch.load(
os.path.join(conf.exp_dir, 'ckpts',
'%d-train_dataset.pth' % start_epoch))
else:
train_dataset.load_data(train_data_list)
utils.printout(conf.flog, str(train_dataset))
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=conf.batch_size, shuffle=True, pin_memory=True, \
num_workers=0, drop_last=True, collate_fn=utils.collate_feats, worker_init_fn=utils.worker_init_fn)
train_num_batch = len(train_dataloader)
### print log
if not conf.no_console_log:
utils.printout(conf.flog, f'training run {conf.exp_name}')
utils.printout(conf.flog, header)
train_batches = enumerate(train_dataloader, 0)
val_batches = enumerate(val_dataloader, 0)
train_fraction_done = 0.0
val_fraction_done = 0.0
val_batch_ind = -1
### train for every batch
for train_batch_ind, batch in train_batches:
train_fraction_done = (train_batch_ind + 1) / train_num_batch
train_step = epoch * train_num_batch + train_batch_ind
log_console = not conf.no_console_log and (last_train_console_log_step is None or \
train_step - last_train_console_log_step >= conf.console_log_interval)
if log_console:
last_train_console_log_step = train_step
# save checkpoint
if train_batch_ind == 0:
with torch.no_grad():
utils.printout(conf.flog, 'Saving checkpoint ...... ')
torch.save(
network.state_dict(),
os.path.join(conf.exp_dir, 'ckpts',
'%d-network.pth' % epoch))
torch.save(
network_opt.state_dict(),
os.path.join(conf.exp_dir, 'ckpts',
'%d-optimizer.pth' % epoch))
torch.save(
network_lr_scheduler.state_dict(),
os.path.join(conf.exp_dir, 'ckpts',
'%d-lr_scheduler.pth' % epoch))
torch.save(
train_dataset,
os.path.join(conf.exp_dir, 'ckpts',
'%d-train_dataset.pth' % epoch))
utils.printout(conf.flog, 'DONE')
# set models to training mode
network.train()
# forward pass (including logging)
total_loss, whole_feats, whole_pcs, whole_pxids, whole_movables = forward(batch=batch, data_features=data_features, network=network, conf=conf, is_val=False, \
step=train_step, epoch=epoch, batch_ind=train_batch_ind, num_batch=train_num_batch, start_time=start_time, \
log_console=log_console, log_tb=not conf.no_tb_log, tb_writer=train_writer, lr=network_opt.param_groups[0]['lr'])
# optimize one step
network_opt.zero_grad()
total_loss.backward()
network_opt.step()
network_lr_scheduler.step()
# sample succ
if conf.sample_succ:
network.eval()
with torch.no_grad():
# sample a random EE orientation
random_up = torch.randn(conf.batch_size,
3).float().to(conf.device)
random_forward = torch.randn(conf.batch_size,
3).float().to(conf.device)
random_left = torch.cross(random_up, random_forward)
random_forward = torch.cross(random_left, random_up)
random_dirs1 = F.normalize(random_up, dim=1).float()
random_dirs2 = F.normalize(random_forward, dim=1).float()
# test over the entire image
whole_pc_scores1 = network.inference_whole_pc(
whole_feats, random_dirs1, random_dirs2) # B x N
whole_pc_scores2 = network.inference_whole_pc(
whole_feats, -random_dirs1, random_dirs2) # B x N
# add to the sample_succ_list if wanted
ss_cur_dir = batch[data_features.index('cur_dir')]
ss_shape_id = batch[data_features.index('shape_id')]
ss_trial_id = batch[data_features.index('trial_id')]
ss_is_original = batch[data_features.index('is_original')]
for i in range(conf.batch_size):
valid_id_l = conf.num_interaction_data_offline + conf.num_interaction_data * (
epoch - 1)
valid_id_r = conf.num_interaction_data_offline + conf.num_interaction_data * epoch
if ('sample-succ' not in ss_cur_dir[i]) and (ss_is_original[i]) and (ss_cur_dir[i] not in sample_succ_dirs) \
and (valid_id_l <= int(ss_trial_id[i]) < valid_id_r):
sample_succ_dirs.append(ss_cur_dir[i])
# choose one from the two options
gt_movable = whole_movables[i].cpu().numpy()
whole_pc_score1 = whole_pc_scores1[i].cpu().numpy(
) * gt_movable
whole_pc_score1[whole_pc_score1 < 0.5] = 0
whole_pc_score_sum1 = np.sum(
whole_pc_score1) + 1e-12
whole_pc_score2 = whole_pc_scores2[i].cpu().numpy(
) * gt_movable
whole_pc_score2[whole_pc_score2 < 0.5] = 0
whole_pc_score_sum2 = np.sum(
whole_pc_score2) + 1e-12
choose1or2_ratio = whole_pc_score_sum1 / (
whole_pc_score_sum1 + whole_pc_score_sum2)
random_dir1 = random_dirs1[i].cpu().numpy()
random_dir2 = random_dirs2[i].cpu().numpy()
if np.random.random() < choose1or2_ratio:
whole_pc_score = whole_pc_score1
else:
whole_pc_score = whole_pc_score2
random_dir1 = -random_dir1
# sample <X, Y> on each img
pp = whole_pc_score + 1e-12
ptid = np.random.choice(len(whole_pc_score),
1,
p=pp / pp.sum())
X = whole_pxids[i, ptid, 0].item()
Y = whole_pxids[i, ptid, 1].item()
# add job to the queue
str_cur_dir1 = ',' + ','.join(
['%f' % elem for elem in random_dir1])
str_cur_dir2 = ',' + ','.join(
['%f' % elem for elem in random_dir2])
sample_succ_list.append((conf.offline_data_dir, str_cur_dir1, str_cur_dir2, \
ss_cur_dir[i].split('/')[-1], cur_sample_succ_dir, X, Y))
# validate one batch
while val_fraction_done <= train_fraction_done and val_batch_ind + 1 < val_num_batch:
val_batch_ind, val_batch = next(val_batches)
val_fraction_done = (val_batch_ind + 1) / val_num_batch
val_step = (epoch + val_fraction_done) * train_num_batch - 1
log_console = not conf.no_console_log and (last_val_console_log_step is None or \
val_step - last_val_console_log_step >= conf.console_log_interval)
if log_console:
last_val_console_log_step = val_step
# set models to evaluation mode
network.eval()
with torch.no_grad():
# forward pass (including logging)
__ = forward(batch=val_batch, data_features=data_features, network=network, conf=conf, is_val=True, \
step=val_step, epoch=epoch, batch_ind=val_batch_ind, num_batch=val_num_batch, start_time=start_time, \
log_console=log_console, log_tb=not conf.no_tb_log, tb_writer=val_writer, lr=network_opt.param_groups[0]['lr'])
print(conf.category_types)
cat2freq = dict()
with open(conf.ins_cnt_fn, 'r') as fin:
for l in fin.readlines():
cat, _, freq = l.rstrip().split()
cat2freq[cat] = int(freq)
print(cat2freq)
datagen = DataGen(conf.num_processes)
with open(conf.data_fn, 'r') as fin:
for l in fin.readlines():
shape_id, cat = l.rstrip().split()
if cat in conf.category_types:
for primact_type in conf.primact_types:
for epoch in range(conf.starting_epoch,
conf.starting_epoch + conf.num_epochs):
for cnt_id in range(cat2freq[cat]):
#print(shape_id, cat, epoch, cnt_id)
datagen.add_one_collect_job(conf.data_dir, shape_id,
cat, cnt_id, primact_type,
epoch)
datagen.start_all()
data_tuple_list = datagen.join_all()
with open(os.path.join(conf.data_dir, conf.out_fn), 'w') as fout:
for item in data_tuple_list:
fout.write(item.split('/')[-1] + '\n')