def train(): net.train() current_lr = config['learning_rate'] print('Loading Dataset...') dataset = MOTTrainDataset(args.mot_root, SSJAugmentation( sst_dim, means ) ) epoch_size = len(dataset) // args.batch_size print('Training SSJ on', dataset.dataset_name) step_index = 0 batch_iterator = None data_loader = data.DataLoader(dataset, batch_size, num_workers=args.num_workers, shuffle=True, collate_fn=collate_fn, pin_memory=False) for iteration in range(args.start_iter, max_iter): if (not batch_iterator) or (iteration % epoch_size == 0): # create batch iterator batch_iterator = iter(data_loader) all_epoch_loss = [] if iteration in stepvalues: step_index += 1 current_lr = adjust_learning_rate(optimizer, args.gamma, step_index) # load train data img_pre, img_next, boxes_pre, boxes_next, labels, valid_pre, valid_next=\ next(batch_iterator) if args.cuda: img_pre = Variable(img_pre.cuda()) img_next = Variable(img_next.cuda()) boxes_pre = Variable(boxes_pre.cuda()) boxes_next = Variable(boxes_next.cuda()) valid_pre = Variable(valid_pre.cuda(), volatile=True) valid_next = Variable(valid_next.cuda(), volatile=True) labels = Variable(labels.cuda(), volatile=True) else: img_pre = Variable(img_pre) img_next = Variable(img_next) boxes_pre = Variable(boxes_pre) boxes_next = Variable(boxes_next) valid_pre = Variable(valid_pre) valid_next = Variable(valid_next) labels = Variable(labels, volatile=True) # forward t0 = time.time() out = net(img_pre, img_next, boxes_pre, boxes_next, valid_pre, valid_next) optimizer.zero_grad() loss_pre, loss_next, loss_similarity, loss, accuracy_pre, accuracy_next, accuracy, predict_indexes = criterion(out, labels, valid_pre, valid_next) loss.backward() optimizer.step() t1 = time.time() all_epoch_loss += [loss.data.cpu()] if iteration % 10 == 0: print('Timer: %.4f sec.' % (t1 - t0)) print('iter ' + repr(iteration) + ', ' + repr(epoch_size) + ' || epoch: %.4f ' % (iteration/(float)(epoch_size)) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ') if args.tensorboard: if len(all_epoch_loss) > 30: writer.add_scalar('data/epoch_loss', float(np.mean(all_epoch_loss)), iteration) writer.add_scalar('data/learning_rate', current_lr, iteration) writer.add_scalar('loss/loss', loss.data.cpu(), iteration) writer.add_scalar('loss/loss_pre', loss_pre.data.cpu(), iteration) writer.add_scalar('loss/loss_next', loss_next.data.cpu(), iteration) writer.add_scalar('loss/loss_similarity', loss_similarity.data.cpu(), iteration) writer.add_scalar('accuracy/accuracy', accuracy.data.cpu(), iteration) writer.add_scalar('accuracy/accuracy_pre', accuracy_pre.data.cpu(), iteration) writer.add_scalar('accuracy/accuracy_next', accuracy_next.data.cpu(), iteration) # add weights if iteration % 1000 == 0: for name, param in net.named_parameters(): writer.add_histogram(name, param.clone().cpu().data.numpy(), iteration) # add images if args.send_images and iteration % 1000 == 0: result_image = show_batch_circle_image(img_pre, img_next, boxes_pre, boxes_next, valid_pre, valid_next, predict_indexes, iteration) writer.add_image('WithLabel/ImageResult', vutils.make_grid(result_image, nrow=2, normalize=True, scale_each=True), iteration) if iteration % save_weights_iteration == 0: print('Saving state, iter:', iteration) torch.save(sst_net.state_dict(), os.path.join( args.save_folder, 'sst300_0712_' + repr(iteration) + '.pth')) torch.save(sst_net.state_dict(), args.save_folder + '' + args.version + '.pth')
continue if config['cuda']: img_pre = Variable(img_pre.cuda()) img_next = Variable(img_next.cuda()) boxes_pre = Variable(boxes_pre.cuda()) boxes_next = Variable(boxes_next.cuda()) valid_pre = Variable(valid_pre.cuda(), volatile=True) valid_next = Variable(valid_next.cuda(), volatile=True) labels = Variable(labels.cuda(), volatile=True) else: img_pre = Variable(img_pre) img_next = Variable(img_next) boxes_pre = Variable(boxes_pre) boxes_next = Variable(boxes_next) valid_pre = Variable(valid_pre) valid_next = Variable(valid_next) labels = Variable(labels, volatile=True) out = sst(img_pre, img_next, boxes_pre, boxes_next, valid_pre, valid_next) loss_pre, loss_next, loss_similarity, loss, accuracy_pre, accuracy_next, accuracy, predict_indexes = criterion( out, labels, valid_pre, valid_next) result_image = show_batch_circle_image(img_pre, img_next, boxes_pre, boxes_next, valid_pre, valid_next, predict_indexes, i) result_image = result_image[0, :].permute(1, 2, 0).cpu().numpy() cv2.imwrite(image_format.format(i), result_image) print(i)
def train(): net.train() current_lr = config['learning_rate'] print('Loading Dataset...') dataset = TrainDataset( args.nuscenes_data_root, SSJAugmentation( config['image_size'], means, ), ) print('length = ', len(dataset)) epoch_size = len(dataset) // args.batch_size step_index = 0 batch_iterator = None batch_size = 1 data_loader = data.DataLoader( dataset, batch_size, num_workers=args.num_workers, shuffle=True, collate_fn=collate_fn, pin_memory=False, ) if args.Joint: nusc = NuScenes( version='v1.0-trainval', verbose=True, dataroot=args.nuscenes_root, ) for iteration in range(0, 380000): print(iteration) if (not batch_iterator) or (iteration % epoch_size == 0): # create batch iterator batch_iterator = iter(data_loader) all_epoch_loss = [] print('pass 0') if iteration in stepvalues: step_index += 1 current_lr = current_lr * 0.5 for param_group in optimizer.param_groups: param_group['lr'] = current_lr # load train data img_pre, img_next, boxes_pre, boxes_next, labels, valid_pre, valid_next, current_poses, next_poses, pre_bbox, next_bbox, img_org_pre, img_org_next, current_tokens, next_tokens =\ next(batch_iterator) if args.Joint: current_token = current_tokens[0] next_token = next_tokens[0] first_token = current_tokens[1] current_cam_record = nusc.get('sample_data', current_token) current_cam_path = nusc.get_sample_data_path(current_token) current_cam_path, boxes, current_camera_intrinsic = nusc.get_sample_data( current_cam_record['token'], ) current_cs_record = nusc.get( 'calibrated_sensor', current_cam_record['calibrated_sensor_token'], ) current_poserecord = nusc.get( 'ego_pose', current_cam_record['ego_pose_token'], ) next_cam_record = nusc.get('sample_data', next_token) next_cam_path = nusc.get_sample_data_path(next_token) next_cam_path, boxes, next_camera_intrinsic = nusc.get_sample_data( next_cam_record['token'], ) next_cs_record = nusc.get( 'calibrated_sensor', next_cam_record['calibrated_sensor_token'], ) next_poserecord = nusc.get( 'ego_pose', next_cam_record['ego_pose_token'], ) first_cam_record = nusc.get('sample_data', first_token) first_cam_path = nusc.get_sample_data_path(first_token) first_cam_path, boxes, first_camera_intrinsic = nusc.get_sample_data( first_cam_record['token'], ) first_cs_record = nusc.get( 'calibrated_sensor', first_cam_record['calibrated_sensor_token'], ) first_poserecord = nusc.get( 'ego_pose', first_cam_record['ego_pose_token'], ) if args.cuda: img_pre = (img_pre.cuda()) img_next = (img_next.cuda()) boxes_pre = (boxes_pre.cuda()) boxes_next = (boxes_next.cuda()) valid_pre = (valid_pre.cuda()) valid_next = (valid_next.cuda()) labels = (labels.cuda()) current_poses = (current_poses.cuda()) next_poses = (next_poses.cuda()) pre_bbox = (pre_bbox.cuda()) next_bbox = (next_bbox.cuda()) img_org_pre = (img_org_pre.cuda()) img_org_next = (img_org_next.cuda()) # forward if args.Joint: out, pose_loss = net( args.Joint, img_pre, img_next, boxes_pre, boxes_next, valid_pre, valid_next, current_poses, next_poses, pre_bbox, next_bbox, img_org_pre, img_org_next, current_cs_record, current_poserecord, next_cs_record, next_poserecord, first_cs_record, first_poserecord, first_camera_intrinsic, ) else: out = net( img_pre, img_next, boxes_pre, boxes_next, valid_pre, valid_next, ) optimizer.zero_grad() loss_pre, loss_next, loss_similarity, loss, accuracy_pre, accuracy_next, accuracy, predict_indexes = criterion( out, labels, valid_pre, valid_next, ) total_loss = loss + 0.1 * pose_loss total_loss.backward() optimizer.step() optimizer.zero_grad() all_epoch_loss += [loss.data.cpu()] print( 'iter ' + repr(iteration) + ', ' + repr(epoch_size) + ' || epoch: %.4f ' % (iteration / (float)(epoch_size)) + ' || Loss: %.4f ||' % (loss.data.item()), end=' ', ) if iteration % 1000 == 0: result_image = show_batch_circle_image( img_pre, img_next, boxes_pre, boxes_next, valid_pre, valid_next, predict_indexes, iteration, ) if iteration % 1000 == 0: print('Saving state, iter:', iteration) torch.save( net.state_dict(), os.path.join( args.save_folder, 'model_' + repr(iteration) + '.pth', ), )