def main_worker(gpu_idx, configs):
configs.gpu_idx = gpu_idx
if configs.gpu_idx is not None:
print("Use GPU: {} for training".format(configs.gpu_idx))
configs.device = torch.device('cuda:{}'.format(configs.gpu_idx))
if configs.distributed:
if configs.dist_url == "env://" and configs.rank == -1:
configs.rank = int(os.environ["RANK"])
if configs.multiprocessing_distributed:
configs.rank = configs.rank * configs.ngpus_per_node + gpu_idx
dist.init_process_group(backend=configs.dist_backend,
init_method=configs.dist_url,
world_size=configs.world_size,
rank=configs.rank)
configs.is_master_node = (not configs.distributed) or (
configs.distributed and (configs.rank % configs.ngpus_per_node == 0))
# model
model = get_model(configs)
model = make_data_parallel(model, configs)
if configs.is_master_node:
num_parameters = get_num_parameters(model)
print('number of trained parameters of the model: {}'.format(
num_parameters))
if configs.pretrained_path is not None:
model = load_pretrained_model(model, configs.pretrained_path, gpu_idx,
configs.overwrite_global_2_local)
# Load dataset
test_loader = create_test_dataloader(configs)
test(test_loader, model, configs)
def demo(configs):
video_loader = TTNet_Video_Loader(configs.video_path, configs.input_size, configs.num_frames_sequence)
result_filename = os.path.join(configs.save_demo_dir, 'results.txt')
frame_rate = video_loader.video_fps
if configs.save_demo_output:
configs.frame_dir = os.path.join(configs.save_demo_dir, 'frame')
if not os.path.isdir(configs.frame_dir):
os.makedirs(configs.frame_dir)
configs.device = torch.device('cuda:{}'.format(configs.gpu_idx))
# model
model = get_model(configs)
model.cuda()
assert configs.pretrained_path is not None, "Need to load the pre-trained model"
model = load_pretrained_model(model, configs.pretrained_path, configs.gpu_idx, configs.overwrite_global_2_local)
model.eval()
middle_idx = int(configs.num_frames_sequence / 2)
queue_frames = deque(maxlen=middle_idx + 1)
frame_idx = 0
w_original, h_original = 1920, 1080
w_resize, h_resize = 320, 128
w_ratio = w_original / w_resize
h_ratio = h_original / h_resize
with torch.no_grad():
for count, origin_imgs, resized_imgs in video_loader:
# take the middle one
img = np.copy(origin_imgs[3 * middle_idx: 3 * (middle_idx + 1), :, :]).transpose(1, 2, 0)
# Expand the first dim
resized_imgs = torch.from_numpy(resized_imgs).to(configs.device, non_blocking=True).float().unsqueeze(0)
origin_imgs = torch.from_numpy(origin_imgs).to(configs.device, non_blocking=True).float().unsqueeze(0)
pred_ball_global, pred_ball_local, pred_events, pred_seg = model.run_demo(origin_imgs, resized_imgs)
prediction_global, prediction_local, prediction_seg, prediction_events = post_processing(
pred_ball_global, pred_ball_local, pred_events, pred_seg, configs.input_size[0],
configs.thresh_ball_pos_mask, configs.seg_thresh, configs.event_thresh)
prediction_ball_final = [
int(prediction_global[0] * w_ratio + prediction_local[0] - w_resize / 2),
int(prediction_global[1] * h_ratio + prediction_local[1] - h_resize / 2)
]
# Get infor of the (middle_idx + 1)th frame
if len(queue_frames) == middle_idx + 1:
frame_pred_infor = queue_frames.popleft()
seg_img = frame_pred_infor['seg'].astype(np.uint8)
ball_pos = frame_pred_infor['ball']
seg_img = cv2.resize(seg_img, (w_original, h_original))
ploted_img = plot_detection(img, ball_pos, seg_img, prediction_events)
ploted_img = cv2.cvtColor(ploted_img, cv2.COLOR_RGB2BGR)
if configs.show_image:
cv2.imshow('ploted_img', ploted_img)
cv2.waitKey(10)
if configs.save_demo_output:
cv2.imwrite(os.path.join(configs.frame_dir, '{:06d}.jpg'.format(frame_idx)), ploted_img)
frame_pred_infor = {
'seg': prediction_seg,
'ball': prediction_ball_final
}
queue_frames.append(frame_pred_infor)
frame_idx += 1
if configs.output_format == 'video':
output_video_path = os.path.join(configs.save_demo_dir, 'result.mp4')
cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -b 5000k -c:v mpeg4 {}'.format(
os.path.join(configs.frame_dir), output_video_path)
os.system(cmd_str)
def main_worker(gpu_idx, configs):
configs.gpu_idx = gpu_idx
if configs.gpu_idx is not None:
print("Use GPU: {} for training".format(configs.gpu_idx))
configs.device = torch.device('cuda:{}'.format(configs.gpu_idx))
if configs.distributed:
if configs.dist_url == "env://" and configs.rank == -1:
configs.rank = int(os.environ["RANK"])
if configs.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
configs.rank = configs.rank * configs.ngpus_per_node + gpu_idx
dist.init_process_group(backend=configs.dist_backend,
init_method=configs.dist_url,
world_size=configs.world_size,
rank=configs.rank)
configs.is_master_node = (not configs.distributed) or (
configs.distributed and (configs.rank % configs.ngpus_per_node == 0))
if configs.is_master_node:
logger = Logger(configs.logs_dir, configs.saved_fn)
logger.info('>>> Created a new logger')
logger.info('>>> configs: {}'.format(configs))
tb_writer = SummaryWriter(
log_dir=os.path.join(configs.logs_dir, 'tensorboard'))
else:
logger = None
tb_writer = None
# model
model = get_model(configs)
# Data Parallel
model = make_data_parallel(model, configs)
# Freeze model
model = freeze_model(model, configs.freeze_modules_list)
if configs.is_master_node:
num_parameters = get_num_parameters(model)
logger.info('number of trained parameters of the model: {}'.format(
num_parameters))
optimizer = get_optimizer(configs, model, is_warm_up=False)
lr_scheduler = get_lr_scheduler(optimizer, configs)
best_val_loss = np.inf
earlystop_count = 0
# optionally load weight from a checkpoint
if configs.pretrained_path is not None:
model = load_pretrained_model(model, configs.pretrained_path, gpu_idx,
configs.overwrite_global_2_local)
if logger is not None:
logger.info('loaded pretrained model at {}'.format(
configs.pretrained_path))
# optionally resume from a checkpoint
if configs.resume_path is not None:
checkpoint = resume_model(configs.resume_path, configs.arch,
configs.gpu_idx)
if hasattr(model, 'module'):
model.module.load_state_dict(checkpoint['state_dict'])
else:
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
best_val_loss = checkpoint['best_val_loss']
earlystop_count = checkpoint['earlystop_count']
configs.start_epoch = checkpoint['epoch'] + 1
if logger is not None:
logger.info(">>> Loading dataset & getting dataloader...")
# Create dataloader
train_loader, val_loader, train_sampler = create_train_val_dataloader(
configs)
if logger is not None:
logger.info('number of batches in train set: {}'.format(
len(train_loader)))
if val_loader is not None:
logger.info('number of batches in val set: {}'.format(
len(val_loader)))
if configs.evaluate:
assert val_loader is not None, "The validation should not be None"
val_loss = validate_one_epoch(val_loader, model,
configs.start_epoch - 1, configs, logger)
print('Evaluate, val_loss: {}'.format(val_loss))
return
for epoch in range(configs.start_epoch, configs.num_epochs + 1):
# Get the current learning rate
for param_group in optimizer.param_groups:
lr = param_group['lr']
if logger is not None:
logger.info('{}'.format('*-' * 40))
logger.info('{} {}/{} {}'.format('=' * 35, epoch,
configs.num_epochs, '=' * 35))
logger.info('{}'.format('*-' * 40))
logger.info('>>> Epoch: [{}/{}] learning rate: {:.2e}'.format(
epoch, configs.num_epochs, lr))
if configs.distributed:
train_sampler.set_epoch(epoch)
# train for one epoch
train_loss = train_one_epoch(train_loader, model, optimizer, epoch,
configs, logger)
# evaluate on validation set
if not configs.no_val:
val_loss = validate_one_epoch(val_loader, model, epoch, configs,
logger)
# Adjust learning rate
if configs.lr_type == 'step_lr':
lr_scheduler.step()
elif configs.lr_type == 'plateau':
assert configs.no_val == True, "Only use plateau when having validation set"
lr_scheduler.step(val_loss)
if not configs.no_val:
is_best = val_loss <= best_val_loss
best_val_loss = min(val_loss, best_val_loss)
print_string = '\t--- train_loss: {:.4f}, val_loss: {:.4f}, best_val_loss: {:.4f}\t'.format(
train_loss, val_loss, best_val_loss)
if tb_writer is not None:
tb_writer.add_scalars('Loss', {
'train': train_loss,
'val': val_loss
}, epoch)
if configs.is_master_node and (is_best or (
(epoch % configs.checkpoint_freq) == 0)):
saved_state = get_saved_state(model, optimizer, lr_scheduler,
epoch, configs, best_val_loss,
earlystop_count)
save_checkpoint(configs.checkpoints_dir, configs.saved_fn,
saved_state, is_best, epoch)
if configs.earlystop_patience:
earlystop_count = 0 if is_best else (earlystop_count + 1)
print_string += ' |||\t earlystop_count: {}'.format(
earlystop_count)
if configs.earlystop_patience <= earlystop_count:
print_string += '\n\t--- Early stopping!!!'
break
else:
print_string += '\n\t--- Continue training..., earlystop_count: {}'.format(
earlystop_count)
if logger is not None:
logger.info(print_string)
else:
if tb_writer is not None:
tb_writer.add_scalars('Loss', {'train': train_loss}, epoch)
if configs.is_master_node and ((epoch % configs.checkpoint_freq)
== 0):
saved_state = get_saved_state(model, optimizer, lr_scheduler,
epoch, configs, best_val_loss,
earlystop_count)
save_checkpoint(configs.checkpoints_dir, configs.saved_fn,
saved_state, False, epoch)
if tb_writer is not None:
tb_writer.close()
cleanup()