def __init__(self, initial_frame, first_box, opts, args, positive=True): self.opts = opts self.positive = positive if positive: self.max_num_past_frames = opts['nFrames_long'] else: self.max_num_past_frames = opts['nFrames_short'] self.transform = ADNet_Augmentation(opts) # list of train_db_ # train_db_ = { # 'past_frame': [], # 'bboxes': [], # 'labels': [], # 'score_labels': [] # } self.train_db = [] self.add_frame_then_generate_samples(initial_frame, first_box)
def adnet_train_sl(args, opts): if torch.cuda.is_available(): if args.cuda: torch.set_default_tensor_type('torch.cuda.FloatTensor') if not args.cuda: print( "WARNING: It looks like you have a CUDA device, but aren't " + "using CUDA.\nRun with --cuda for optimal training speed.") torch.set_default_tensor_type('torch.FloatTensor') else: torch.set_default_tensor_type('torch.FloatTensor') if not os.path.exists(args.save_folder): os.mkdir(args.save_folder) if args.visualize: writer = SummaryWriter(log_dir=os.path.join('tensorboardx_log', args.save_file)) train_videos = get_train_videos(opts) opts['num_videos'] = len(train_videos['video_names']) net, domain_specific_nets = adnet(opts=opts, trained_file=args.resume, multidomain=args.multidomain) if args.cuda: net = nn.DataParallel(net) cudnn.benchmark = True net = net.cuda() if args.cuda: optimizer = optim.SGD([ {'params': net.module.base_network.parameters(), 'lr': 1e-4}, {'params': net.module.fc4_5.parameters()}, {'params': net.module.fc6.parameters()}, {'params': net.module.fc7.parameters()}], # as action dynamic is zero, it doesn't matter lr=1e-3, momentum=opts['train']['momentum'], weight_decay=opts['train']['weightDecay']) else: optimizer = optim.SGD([ {'params': net.base_network.parameters(), 'lr': 1e-4}, {'params': net.fc4_5.parameters()}, {'params': net.fc6.parameters()}, {'params': net.fc7.parameters()}], lr=1e-3, momentum=opts['train']['momentum'], weight_decay=opts['train']['weightDecay']) if args.resume: # net.load_weights(args.resume) checkpoint = torch.load(args.resume) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) net.train() if not args.resume: print('Initializing weights...') if args.cuda: scal = torch.Tensor([0.01]) # fc 4 nn.init.normal_(net.module.fc4_5[0].weight.data) net.module.fc4_5[0].weight.data = net.module.fc4_5[0].weight.data * scal.expand_as(net.module.fc4_5[0].weight.data) net.module.fc4_5[0].bias.data.fill_(0.1) # fc 5 nn.init.normal_(net.module.fc4_5[3].weight.data) net.module.fc4_5[3].weight.data = net.module.fc4_5[3].weight.data * scal.expand_as(net.module.fc4_5[3].weight.data) net.module.fc4_5[3].bias.data.fill_(0.1) # fc 6 nn.init.normal_(net.module.fc6.weight.data) net.module.fc6.weight.data = net.module.fc6.weight.data * scal.expand_as(net.module.fc6.weight.data) net.module.fc6.bias.data.fill_(0) # fc 7 nn.init.normal_(net.module.fc7.weight.data) net.module.fc7.weight.data = net.module.fc7.weight.data * scal.expand_as(net.module.fc7.weight.data) net.module.fc7.bias.data.fill_(0) else: scal = torch.Tensor([0.01]) # fc 4 nn.init.normal_(net.fc4_5[0].weight.data) net.fc4_5[0].weight.data = net.fc4_5[0].weight.data * scal.expand_as(net.fc4_5[0].weight.data ) net.fc4_5[0].bias.data.fill_(0.1) # fc 5 nn.init.normal_(net.fc4_5[3].weight.data) net.fc4_5[3].weight.data = net.fc4_5[3].weight.data * scal.expand_as(net.fc4_5[3].weight.data) net.fc4_5[3].bias.data.fill_(0.1) # fc 6 nn.init.normal_(net.fc6.weight.data) net.fc6.weight.data = net.fc6.weight.data * scal.expand_as(net.fc6.weight.data) net.fc6.bias.data.fill_(0) # fc 7 nn.init.normal_(net.fc7.weight.data) net.fc7.weight.data = net.fc7.weight.data * scal.expand_as(net.fc7.weight.data) net.fc7.bias.data.fill_(0) action_criterion = nn.CrossEntropyLoss() score_criterion = nn.CrossEntropyLoss() print('generating Supervised Learning dataset..') # dataset = SLDataset(train_videos, opts, transform= datasets_pos, datasets_neg = initialize_pos_neg_dataset(train_videos, opts, transform=ADNet_Augmentation(opts)) number_domain = opts['num_videos'] batch_iterators_pos = [] batch_iterators_neg = [] # calculating number of data len_dataset_pos = 0 len_dataset_neg = 0 for dataset_pos in datasets_pos: len_dataset_pos += len(dataset_pos) for dataset_neg in datasets_neg: len_dataset_neg += len(dataset_neg) epoch_size_pos = len_dataset_pos // opts['minibatch_size'] epoch_size_neg = len_dataset_neg // opts['minibatch_size'] epoch_size = epoch_size_pos + epoch_size_neg # 1 epoch, how many iterations print("1 epoch = " + str(epoch_size) + " iterations") max_iter = opts['numEpoch'] * epoch_size print("maximum iteration = " + str(max_iter)) data_loaders_pos = [] data_loaders_neg = [] for dataset_pos in datasets_pos: data_loaders_pos.append(data.DataLoader(dataset_pos, opts['minibatch_size'], num_workers=args.num_workers, shuffle=True, pin_memory=True)) for dataset_neg in datasets_neg: data_loaders_neg.append(data.DataLoader(dataset_neg, opts['minibatch_size'], num_workers=args.num_workers, shuffle=True, pin_memory=True)) epoch = args.start_epoch if epoch != 0 and args.start_iter == 0: start_iter = epoch * epoch_size else: start_iter = args.start_iter which_dataset = list(np.full(epoch_size_pos, fill_value=1)) which_dataset.extend(np.zeros(epoch_size_neg, dtype=int)) shuffle(which_dataset) which_domain = np.random.permutation(number_domain) action_loss = 0 score_loss = 0 # training loop for iteration in range(start_iter, max_iter): if args.multidomain: curr_domain = which_domain[iteration % len(which_domain)] else: curr_domain = 0 # if new epoch (not including the very first iteration) if (iteration != start_iter) and (iteration % epoch_size == 0): epoch += 1 shuffle(which_dataset) np.random.shuffle(which_domain) print('Saving state, epoch:', epoch) domain_specific_nets_state_dict = [] for domain_specific_net in domain_specific_nets: domain_specific_nets_state_dict.append(domain_specific_net.state_dict()) torch.save({ 'epoch': epoch, 'adnet_state_dict': net.state_dict(), 'adnet_domain_specific_state_dict': domain_specific_nets, 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(args.save_folder, args.save_file) + 'epoch' + repr(epoch) + '.pth') if args.visualize: writer.add_scalars('data/epoch_loss', {'action_loss': action_loss / epoch_size, 'score_loss': score_loss / epoch_size, 'total': (action_loss + score_loss) / epoch_size}, global_step=epoch) # reset epoch loss counters action_loss = 0 score_loss = 0 # if new epoch (including the first iteration), initialize the batch iterator # or just resuming where batch_iterator_pos and neg haven't been initialized if iteration % epoch_size == 0 or len(batch_iterators_pos) == 0 or len(batch_iterators_neg) == 0: # create batch iterator for data_loader_pos in data_loaders_pos: batch_iterators_pos.append(iter(data_loader_pos)) for data_loader_neg in data_loaders_neg: batch_iterators_neg.append(iter(data_loader_neg)) # if not batch_iterators_pos[curr_domain]: # # create batch iterator # batch_iterators_pos[curr_domain] = iter(data_loaders_pos[curr_domain]) # # if not batch_iterators_neg[curr_domain]: # # create batch iterator # batch_iterators_neg[curr_domain] = iter(data_loaders_neg[curr_domain]) # load train data if which_dataset[iteration % len(which_dataset)]: # if positive try: images, bbox, action_label, score_label, vid_idx = next(batch_iterators_pos[curr_domain]) except StopIteration: batch_iterators_pos[curr_domain] = iter(data_loaders_pos[curr_domain]) images, bbox, action_label, score_label, vid_idx = next(batch_iterators_pos[curr_domain]) else: try: images, bbox, action_label, score_label, vid_idx = next(batch_iterators_neg[curr_domain]) except StopIteration: batch_iterators_neg[curr_domain] = iter(data_loaders_neg[curr_domain]) images, bbox, action_label, score_label, vid_idx = next(batch_iterators_neg[curr_domain]) # TODO: check if this requires grad is really false like in Variable if args.cuda: images = torch.Tensor(images.cuda()) bbox = torch.Tensor(bbox.cuda()) action_label = torch.Tensor(action_label.cuda()) score_label = torch.Tensor(score_label.float().cuda()) else: images = torch.Tensor(images) bbox = torch.Tensor(bbox) action_label = torch.Tensor(action_label) score_label = torch.Tensor(score_label) t0 = time.time() # load ADNetDomainSpecific with video index if args.cuda: net.module.load_domain_specific(domain_specific_nets[curr_domain]) else: net.load_domain_specific(domain_specific_nets[curr_domain]) # forward action_out, score_out = net(images) # backprop optimizer.zero_grad() if which_dataset[iteration % len(which_dataset)]: # if positive action_l = action_criterion(action_out, torch.max(action_label, 1)[1]) else: action_l = torch.Tensor([0]) score_l = score_criterion(score_out, score_label.long()) loss = action_l + score_l loss.backward() optimizer.step() action_loss += action_l.item() score_loss += score_l.item() # save the ADNetDomainSpecific back to their module if args.cuda: domain_specific_nets[curr_domain].load_weights_from_adnet(net.module) else: domain_specific_nets[curr_domain].load_weights_from_adnet(net) t1 = time.time() if iteration % 10 == 0: print('Timer: %.4f sec.' % (t1 - t0)) print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data.item()), end=' ') if args.visualize and args.send_images_to_visualization: random_batch_index = np.random.randint(images.size(0)) writer.add_image('image', images.data[random_batch_index].cpu().numpy(), random_batch_index) if args.visualize: writer.add_scalars('data/iter_loss', {'action_loss': action_l.item(), 'score_loss': score_l.item(), 'total': (action_l.item() + score_l.item())}, global_step=iteration) # hacky fencepost solution for 0th epoch plot if iteration == 0: writer.add_scalars('data/epoch_loss', {'action_loss': action_loss, 'score_loss': score_loss, 'total': (action_loss + score_loss)}, global_step=epoch) if iteration % 5000 == 0: print('Saving state, iter:', iteration) domain_specific_nets_state_dict = [] for domain_specific_net in domain_specific_nets: domain_specific_nets_state_dict.append(domain_specific_net.state_dict()) torch.save({ 'epoch': epoch, 'adnet_state_dict': net.state_dict(), 'adnet_domain_specific_state_dict': domain_specific_nets, 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(args.save_folder, args.save_file) + repr(iteration) + '_epoch' + repr(epoch) +'.pth') # final save torch.save({ 'epoch': epoch, 'adnet_state_dict': net.state_dict(), 'adnet_domain_specific_state_dict': domain_specific_nets, 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(args.save_folder, args.save_file) + '.pth') return net, domain_specific_nets, train_videos
def reset(self, net, domain_specific_nets, train_videos, opts, args): self.action_list = [] # a_t,l # argmax of self.action_prob_list self.action_prob_list = [] # output of network (fc6_out) self.log_probs_list = [ ] # log probs from each self.action_prob_list member self.reward_list = [] # tracking score self.patch_list = [] # input of network self.action_dynamic_list = [ ] # action_dynamic used for inference (means before updating the action_dynamic) self.result_box_list = [] self.vid_idx_list = [] print('generating reinforcement learning dataset') transform = ADNet_Augmentation(opts) self.env = TrackingEnvironment(train_videos, opts, transform=transform, args=args) clip_idx = 0 while True: # for every clip (l) num_step_history = [] # T_l num_frame = 1 # the first frame won't be tracked.. t = 0 box_history_clip = [] # for checking oscillation in a clip if args.cuda: net.module.reset_action_dynamic() else: net.reset_action_dynamic( ) # action dynamic should be in a clip (what makes sense...) while True: # for every frame in a clip (t) tic = time.time() if args.display_images: im_with_bb = display_result(self.env.get_current_img(), self.env.get_state()) cv2.imshow('patch', self.env.get_current_patch_unprocessed()) cv2.waitKey(1) else: im_with_bb = draw_box(self.env.get_current_img(), self.env.get_state()) if args.save_result_images: cv2.imwrite( 'images/' + str(clip_idx) + '-' + str(t) + '.jpg', im_with_bb) curr_patch = self.env.get_current_patch() if args.cuda: curr_patch = curr_patch.cuda() # self.patch_list.append(curr_patch.cpu().data.numpy()) # TODO: saving patch takes cuda memory # TODO: saving action_dynamic takes cuda memory # if args.cuda: # self.action_dynamic_list.append(net.module.get_action_dynamic()) # else: # self.action_dynamic_list.append(net.get_action_dynamic()) curr_patch = curr_patch.unsqueeze( 0) # 1 batch input [1, curr_patch.shape] # load ADNetDomainSpecific with video index if args.multidomain: vid_idx = self.env.get_current_train_vid_idx() else: vid_idx = 0 if args.cuda: net.module.load_domain_specific( domain_specific_nets[vid_idx]) else: net.load_domain_specific(domain_specific_nets[vid_idx]) fc6_out, fc7_out = net.forward(curr_patch, update_action_dynamic=True) if args.cuda: action = np.argmax(fc6_out.detach().cpu().numpy() ) # TODO: really okay to detach? action_prob = fc6_out.detach().cpu().numpy()[0][action] else: action = np.argmax(fc6_out.detach().numpy() ) # TODO: really okay to detach? action_prob = fc6_out.detach().numpy()[0][action] m = Categorical(probs=fc6_out) action_ = m.sample( ) # action and action_ are same value. Only differ in the type (int and tensor) self.log_probs_list.append( m.log_prob(action_).cpu().data.numpy()) self.vid_idx_list.append(vid_idx) self.action_list.append(action) # TODO: saving action_prob_list takes cuda memory # self.action_prob_list.append(action_prob) new_state, reward, done, info = self.env.step(action) # check oscilating if any((np.array(new_state).round() == x).all() for x in np.array(box_history_clip).round()): action = opts['stop_action'] reward, done, finish_epoch = self.env.go_to_next_frame() info['finish_epoch'] = finish_epoch # check if number of action is already too much if t > opts['num_action_step_max']: action = opts['stop_action'] reward, done, finish_epoch = self.env.go_to_next_frame() info['finish_epoch'] = finish_epoch # TODO: saving result_box takes cuda memory # self.result_box_list.append(list(new_state)) box_history_clip.append(list(new_state)) t += 1 if action == opts['stop_action']: num_frame += 1 num_step_history.append(t) t = 0 toc = time.time() - tic print('forward time (clip ' + str(clip_idx) + " - frame " + str(num_frame) + " - t " + str(t) + ") = " + str(toc) + " s") if done: # if finish the clip break tracking_scores_size = np.array(num_step_history).sum() tracking_scores = np.full( tracking_scores_size, reward) # seems no discount factor whatsoever self.reward_list.extend(tracking_scores) # self.reward_list.append(tracking_scores) clip_idx += 1 if info['finish_epoch']: break print('generating reinforcement learning dataset finish')
def adnet_test(net, vid_path, opts, args): if torch.cuda.is_available(): if args.cuda: torch.set_default_tensor_type('torch.cuda.FloatTensor') if not args.cuda: print( "WARNING: It looks like you have a CUDA device, but aren't " + "using CUDA.\nRun with --cuda for optimal training speed.") torch.set_default_tensor_type('torch.FloatTensor') else: torch.set_default_tensor_type('torch.FloatTensor') transform = ADNet_Augmentation(opts) print('Testing sequences in ' + str(vid_path) + '...') t_sum = 0 if args.visualize: writer = SummaryWriter( log_dir=os.path.join('tensorboardx_log', 'online_adapatation_' + args.save_result_npy)) ################################ # Load video sequences ################################ vid_info = {'gt': [], 'img_files': [], 'nframes': 0} vid_info['img_files'] = glob.glob(os.path.join(vid_path, 'color', '*.jpg')) vid_info['img_files'].sort(key=str.lower) gt_path = os.path.join(vid_path, 'groundtruth.txt') if not os.path.exists(gt_path): bboxes = [] t = 0 return bboxes, t_sum # parse gt gtFile = open(gt_path, 'r') gt = gtFile.read().split('\n') for i in range(len(gt)): if gt[i] == '' or gt[i] is None: continue if ',' in gt[i]: separator = ',' elif '\t' in gt[i]: separator = '\t' elif ' ' in gt[i]: separator = ' ' else: separator = ',' gt[i] = gt[i].split(separator) gt[i] = list(map(float, gt[i])) gtFile.close() if len(gt[0]) >= 6: for gtidx in range(len(gt)): if gt[gtidx] == "": continue x = gt[gtidx][0:len(gt[gtidx]):2] y = gt[gtidx][1:len(gt[gtidx]):2] gt[gtidx] = [min(x), min(y), max(x) - min(x), max(y) - min(y)] vid_info['gt'] = gt if vid_info['gt'][-1] == '': # small hack vid_info['gt'] = vid_info['gt'][:-1] vid_info['nframes'] = min(len(vid_info['img_files']), len(vid_info['gt'])) # catch the first box curr_bbox = vid_info['gt'][0] # init containers bboxes = np.zeros(np.array( vid_info['gt']).shape) # tracking result containers ntraining = 0 # setup training if args.cuda: optimizer = optim.SGD([{ 'params': net.module.base_network.parameters(), 'lr': 0 }, { 'params': net.module.fc4_5.parameters() }, { 'params': net.module.fc6.parameters() }, { 'params': net.module.fc7.parameters(), 'lr': 1e-3 }], lr=1e-3, momentum=opts['train']['momentum'], weight_decay=opts['train']['weightDecay']) else: optimizer = optim.SGD([{ 'params': net.base_network.parameters(), 'lr': 0 }, { 'params': net.fc4_5.parameters() }, { 'params': net.fc6.parameters() }, { 'params': net.fc7.parameters(), 'lr': 1e-3 }], lr=1e-3, momentum=opts['train']['momentum'], weight_decay=opts['train']['weightDecay']) action_criterion = nn.CrossEntropyLoss() score_criterion = nn.CrossEntropyLoss() dataset_storage_pos = None dataset_storage_neg = None is_negative = False # is_negative = True if the tracking failed target_score = 0 all_iteration = 0 t = 0 for idx in range(vid_info['nframes']): # for frame_idx, frame_path in enumerate(vid_info['img_files']): frame_idx = idx frame_path = vid_info['img_files'][idx] t0_wholetracking = time.time() frame = cv2.imread(frame_path) # draw box or with display, then save if args.display_images: im_with_bb = display_result(frame, curr_bbox) # draw box and display else: im_with_bb = draw_box(frame, curr_bbox) if args.save_result_images: filename = os.path.join(args.save_result_images, str(frame_idx) + '-' + str(t) + '.jpg') cv2.imwrite(filename, im_with_bb) curr_bbox_old = curr_bbox cont_negatives = 0 if frame_idx > 0: # tracking if args.cuda: net.module.set_phase('test') else: net.set_phase('test') t = 0 while True: curr_patch, curr_bbox, _, _ = transform( frame, curr_bbox, None, None) if args.cuda: curr_patch = curr_patch.cuda() curr_patch = curr_patch.unsqueeze( 0) # 1 batch input [1, curr_patch.shape] fc6_out, fc7_out = net.forward(curr_patch) curr_score = fc7_out.detach().cpu().numpy()[0][1] if ntraining > args.believe_score_result: if curr_score < opts['failedThre']: cont_negatives += 1 if args.cuda: action = np.argmax(fc6_out.detach().cpu().numpy() ) # TODO: really okay to detach? action_prob = fc6_out.detach().cpu().numpy()[0][action] else: action = np.argmax(fc6_out.detach().numpy() ) # TODO: really okay to detach? action_prob = fc6_out.detach().numpy()[0][action] # do action curr_bbox = do_action(curr_bbox, opts, action, frame.shape) # bound the curr_bbox size if curr_bbox[2] < 10: curr_bbox[0] = min( 0, curr_bbox[0] + curr_bbox[2] / 2 - 10 / 2) curr_bbox[2] = 10 if curr_bbox[3] < 10: curr_bbox[1] = min( 0, curr_bbox[1] + curr_bbox[3] / 2 - 10 / 2) curr_bbox[3] = 10 t += 1 # draw box or with display, then save if args.display_images: im_with_bb = display_result( frame, curr_bbox) # draw box and display else: im_with_bb = draw_box(frame, curr_bbox) if args.save_result_images: filename = os.path.join( args.save_result_images, str(frame_idx) + '-' + str(t) + '.jpg') cv2.imwrite(filename, im_with_bb) if action == opts[ 'stop_action'] or t >= opts['num_action_step_max']: break print('final curr_score: %.4f' % curr_score) # redetection when confidence < threshold 0.5. But when fc7 is already reliable. Else, just trust the ADNet if ntraining > args.believe_score_result: if curr_score < 0.5: print('redetection') is_negative = True # redetection process redet_samples = gen_samples( 'gaussian', curr_bbox_old, opts['redet_samples'], opts, min(1.5, 0.6 * 1.15**cont_negatives), opts['redet_scale_factor']) score_samples = [] for redet_sample in redet_samples: temp_patch, temp_bbox, _, _ = transform( frame, redet_sample, None, None) if args.cuda: temp_patch = temp_patch.cuda() temp_patch = temp_patch.unsqueeze( 0) # 1 batch input [1, curr_patch.shape] fc6_out_temp, fc7_out_temp = net.forward(temp_patch) score_samples.append( fc7_out_temp.detach().cpu().numpy()[0][1]) score_samples = np.array(score_samples) max_score_samples_idx = np.argmax(score_samples) # replace the curr_box with the samples with maximum score curr_bbox = redet_samples[max_score_samples_idx] # update the final result image if args.display_images: im_with_bb = display_result( frame, curr_bbox) # draw box and display else: im_with_bb = draw_box(frame, curr_bbox) if args.save_result_images: filename = os.path.join(args.save_result_images, str(frame_idx) + '-redet.jpg') cv2.imwrite(filename, im_with_bb) else: is_negative = False else: is_negative = False if args.save_result_images: filename = os.path.join(args.save_result_images, 'final-' + str(frame_idx) + '.jpg') cv2.imwrite(filename, im_with_bb) # record the curr_bbox result bboxes[frame_idx] = curr_bbox # create or update storage + set iteration_range for training if frame_idx == 0: dataset_storage_pos = OnlineAdaptationDatasetStorage( initial_frame=frame, first_box=curr_bbox, opts=opts, args=args, positive=True) if opts['nNeg_init'] != 0: # (thanks to small hack in adnet_test) the nNeg_online is also 0 dataset_storage_neg = OnlineAdaptationDatasetStorage( initial_frame=frame, first_box=curr_bbox, opts=opts, args=args, positive=False) iteration_range = range(opts['finetune_iters']) else: assert dataset_storage_pos is not None if opts['nNeg_init'] != 0: # (thanks to small hack in adnet_test) the nNeg_online is also 0 assert dataset_storage_neg is not None # if confident or when always generate samples, generate new samples if ntraining < args.believe_score_result: always_generate_samples = True # as FC7 wasn't trained, it is better to wait for some time to believe its confidence result to decide whether to generate samples or not.. Before believe it, better to just generate sample always else: always_generate_samples = False if always_generate_samples or (not is_negative or target_score > opts['successThre']): dataset_storage_pos.add_frame_then_generate_samples( frame, curr_bbox) iteration_range = range(opts['finetune_iters_online']) # training when depend on the frequency.. else, don't run the training code... if False and frame_idx % args.online_adaptation_every_I_frames == 0: ntraining += 1 # generate dataset just before training dataset_pos = OnlineAdaptationDataset(dataset_storage_pos) data_loader_pos = data.DataLoader(dataset_pos, opts['minibatch_size'], num_workers=args.num_workers, shuffle=True, pin_memory=False) batch_iterator_pos = None if opts['nNeg_init'] != 0: # (thanks to small hack in adnet_test) the nNeg_online is also 0 dataset_neg = OnlineAdaptationDataset(dataset_storage_neg) data_loader_neg = data.DataLoader(dataset_neg, opts['minibatch_size'], num_workers=args.num_workers, shuffle=True, pin_memory=False) batch_iterator_neg = None else: dataset_neg = [] epoch_size_pos = len(dataset_pos) // opts['minibatch_size'] epoch_size_neg = len(dataset_neg) // opts['minibatch_size'] epoch_size = epoch_size_pos + epoch_size_neg # 1 epoch, how many iterations which_dataset = list(np.full(epoch_size_pos, fill_value=1)) which_dataset.extend(np.zeros(epoch_size_neg, dtype=int)) shuffle(which_dataset) print("1 epoch = " + str(epoch_size) + " iterations") if args.cuda: net.module.set_phase('train') else: net.set_phase('train') # training loop for iteration in iteration_range: all_iteration += 1 # use this for update the visualization # create batch iterator if (not batch_iterator_pos) or (iteration % epoch_size == 0): batch_iterator_pos = iter(data_loader_pos) if opts['nNeg_init'] != 0: if (not batch_iterator_neg) or (iteration % epoch_size == 0): batch_iterator_neg = iter(data_loader_neg) # load train data if which_dataset[iteration % len(which_dataset)]: # if positive images, bbox, action_label, score_label = next( batch_iterator_pos) else: images, bbox, action_label, score_label = next( batch_iterator_neg) if args.cuda: images = torch.Tensor(images.cuda()) bbox = torch.Tensor(bbox.cuda()) action_label = torch.Tensor(action_label.cuda()) score_label = torch.Tensor(score_label.float().cuda()) else: images = torch.Tensor(images) bbox = torch.Tensor(bbox) action_label = torch.Tensor(action_label) score_label = torch.Tensor(score_label) # forward t0 = time.time() action_out, score_out = net(images) # backprop optimizer.zero_grad() if which_dataset[iteration % len(which_dataset)]: # if positive action_l = action_criterion(action_out, torch.max(action_label, 1)[1]) else: action_l = torch.Tensor([0]) score_l = score_criterion(score_out, score_label.long()) loss = action_l + score_l loss.backward() optimizer.step() t1 = time.time() if all_iteration % 10 == 0: print('Timer: %.4f sec.' % (t1 - t0)) print('iter ' + repr(all_iteration) + ' || Loss: %.4f ||' % (loss.data.item()), end=' ') if args.visualize and args.send_images_to_visualization: random_batch_index = np.random.randint(images.size(0)) writer.add_image( 'image', images.data[random_batch_index].cpu().numpy(), random_batch_index) if args.visualize: writer.add_scalars( 'data/iter_loss', { 'action_loss': action_l.item(), 'score_loss': score_l.item(), 'total': (action_l.item() + score_l.item()) }, global_step=all_iteration) t1_wholetracking = time.time() t_sum += t1_wholetracking - t0_wholetracking print('whole tracking time = %.4f sec.' % (t1_wholetracking - t0_wholetracking)) # evaluate the precision bboxes = np.array(bboxes) vid_info['gt'] = np.array(vid_info['gt']) # iou_precisions = iou_precision_plot(bboxes, vid_info['gt'], vid_path, show=args.display_images, save_plot=args.save_result_images) # # distance_precisions = distance_precision_plot(bboxes, vid_info['gt'], vid_path, show=args.display_images, save_plot=args.save_result_images) # # precisions = [distance_precisions, iou_precisions] np.save(args.save_result_npy + '-bboxes.npy', bboxes) np.save(args.save_result_npy + '-ground_truth.npy', vid_info['gt']) # return bboxes, t_sum, precisions return bboxes, t_sum
def adnet_train_sl_mot(args, opts, mot, num_obj_to_track=2): if torch.cuda.is_available(): if args.cuda: torch.set_default_tensor_type('torch.cuda.FloatTensor') if not args.cuda: print( "WARNING: It looks like you have a CUDA device, but aren't " + "using CUDA.\nRun with --cuda for optimal training speed.") torch.set_default_tensor_type('torch.FloatTensor') else: torch.set_default_tensor_type('torch.FloatTensor') if not os.path.exists(args.save_folder): os.mkdir(args.save_folder) if args.visualize: writer = SummaryWriter( log_dir=os.path.join('tensorboardx_log', args.save_file)) train_videos = get_train_videos(opts) opts['num_videos'] = len(train_videos['video_names']) net, domain_specific_nets = adnet_mot(opts=opts, trained_file=args.resume, multidomain=args.multidomain) if args.cuda: net = nn.DataParallel(net) cudnn.benchmark = True net = net.cuda() if args.cuda: optimizer = optim.Adam( [{ 'params': net.module.base_network.parameters(), 'lr': 1e-4 }, { 'params': net.module.fc4_5.parameters() }, { 'params': net.module.fc6.parameters() }, { 'params': net.module.fc7.parameters() }], # as action dynamic is zero, it doesn't matter lr=1e-3, weight_decay=opts['train']['weightDecay']) else: optimizer = optim.SGD([{ 'params': net.base_network.parameters(), 'lr': 1e-4 }, { 'params': net.fc4_5.parameters() }, { 'params': net.fc6.parameters() }, { 'params': net.fc7.parameters() }], lr=1e-3, momentum=opts['train']['momentum'], weight_decay=opts['train']['weightDecay']) if args.resume: # net.load_weights(args.resume) checkpoint = torch.load(args.resume) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) net.train() if not args.resume: print('Initializing weights...') if args.cuda: norm_std = 0.01 # fc 4 nn.init.normal_(net.module.fc4_5[0].weight.data, std=norm_std) net.module.fc4_5[0].bias.data.fill_(0.1) # fc 5 nn.init.normal_(net.module.fc4_5[3].weight.data, std=norm_std) net.module.fc4_5[3].bias.data.fill_(0.1) # fc 6 nn.init.normal_(net.module.fc6.weight.data, std=norm_std) net.module.fc6.bias.data.fill_(0) # fc 7 nn.init.normal_(net.module.fc7.weight.data, std=norm_std) net.module.fc7.bias.data.fill_(0) else: scal = torch.Tensor([0.01]) # fc 4 nn.init.normal_(net.fc4_5[0].weight.data) net.fc4_5[0].weight.data = net.fc4_5[ 0].weight.data * scal.expand_as(net.fc4_5[0].weight.data) net.fc4_5[0].bias.data.fill_(0.1) # fc 5 nn.init.normal_(net.fc4_5[3].weight.data) net.fc4_5[3].weight.data = net.fc4_5[ 3].weight.data * scal.expand_as(net.fc4_5[3].weight.data) net.fc4_5[3].bias.data.fill_(0.1) # fc 6 nn.init.normal_(net.fc6.weight.data) net.fc6.weight.data = net.fc6.weight.data * scal.expand_as( net.fc6.weight.data) net.fc6.bias.data.fill_(0) # fc 7 nn.init.normal_(net.fc7.weight.data) net.fc7.weight.data = net.fc7.weight.data * scal.expand_as( net.fc7.weight.data) net.fc7.bias.data.fill_(0) action_criterion = nn.BCEWithLogitsLoss() score_criterion = nn.BCEWithLogitsLoss() print('generating Supervised Learning dataset..') # dataset = SLDataset(train_videos, opts, transform= datasets_pos, datasets_neg = initialize_pos_neg_dataset_adnet_mot( train_videos, opts, transform=ADNet_Augmentation(opts)) number_domain = opts['num_videos'] assert number_domain == len( datasets_pos ), "Num videos given in opts is incorrect! It should be {}".format( len(datasets_neg)) batch_iterators_pos_train = [] batch_iterators_neg_train = [] action_loss_tr = 0 score_loss_tr = 0 # calculating number of data len_dataset_pos = 0 len_dataset_neg = 0 for dataset_pos in datasets_pos: len_dataset_pos += len(dataset_pos) for dataset_neg in datasets_neg: len_dataset_neg += len(dataset_neg) epoch_size_pos = len_dataset_pos // opts['minibatch_size'] epoch_size_neg = len_dataset_neg // opts['minibatch_size'] epoch_size = epoch_size_pos + epoch_size_neg # 1 epoch, how many iterations print("1 epoch = " + str(epoch_size) + " iterations") max_iter = opts['numEpoch'] * epoch_size print("maximum iteration = " + str(max_iter)) data_loaders_pos_train = [] data_loaders_pos_val = [] data_loaders_neg_train = [] data_loaders_neg_val = [] for dataset_pos in datasets_pos: # num_val = int(opts['val_percent'] * len(dataset_pos)) num_val = 1 num_train = len(dataset_pos) - num_val train, valid = torch.utils.data.random_split(dataset_pos, [num_train, num_val]) data_loaders_pos_train.append( data.DataLoader(train, opts['minibatch_size'], num_workers=2, shuffle=True, pin_memory=True)) data_loaders_pos_val.append( data.DataLoader(valid, opts['minibatch_size'], num_workers=0, shuffle=True, pin_memory=False)) for dataset_neg in datasets_neg: num_val = int(opts['val_percent'] * len(dataset_neg)) num_train = len(dataset_neg) - num_val train, valid = torch.utils.data.random_split(dataset_neg, [num_train, num_val]) data_loaders_neg_train.append( data.DataLoader(train, opts['minibatch_size'], num_workers=1, shuffle=True, pin_memory=True)) data_loaders_neg_val.append( data.DataLoader(valid, opts['minibatch_size'], num_workers=0, shuffle=True, pin_memory=False)) epoch = args.start_epoch if epoch != 0 and args.start_iter == 0: start_iter = epoch * epoch_size else: start_iter = args.start_iter which_dataset = list(np.full(epoch_size_pos, fill_value=1)) which_dataset.extend(np.zeros(epoch_size_neg, dtype=int)) shuffle(which_dataset) which_dataset = torch.Tensor(which_dataset).cuda() which_domain = np.random.permutation(number_domain) # training loop time_arr = np.zeros(10) for iteration in tqdm(range(start_iter, max_iter)): t0 = time.time() if args.multidomain: curr_domain = which_domain[iteration % len(which_domain)] else: curr_domain = 0 # if new epoch (not including the very first iteration) if (iteration != start_iter) and (iteration % epoch_size == 0): epoch += 1 shuffle(which_dataset) np.random.shuffle(which_domain) print('Saving state, epoch: {}'.format(epoch)) domain_specific_nets_state_dict = [] for domain_specific_net in domain_specific_nets: domain_specific_nets_state_dict.append( domain_specific_net.state_dict()) torch.save( { 'epoch': epoch, 'adnet_state_dict': net.state_dict(), 'adnet_domain_specific_state_dict': domain_specific_nets, 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(args.save_folder, args.save_file) + 'epoch' + repr(epoch) + '.pth') # VAL # for curr_domain_temp in range(number_domain): # accuracy = [] # action_loss_val = [] # score_loss_val = [] # # # load ADNetDomainSpecific with video index # if args.cuda: # net.module.load_domain_specific(domain_specific_nets[curr_domain_temp]) # else: # net.load_domain_specific(domain_specific_nets[curr_domain_temp]) # for i, temp in enumerate( # [data_loaders_pos_val[curr_domain_temp], data_loaders_neg_val[curr_domain_temp]]): # for images, bbox, action_label, score_label, _ in temp: # images = images.to('cuda', non_blocking=True) # action_label = action_label.to('cuda', non_blocking=True) # score_label = score_label.float().to('cuda', non_blocking=True) # # # forward # action_out, score_out = net(images) # # if i == 0: # if positive # action_l = action_criterion(action_out, torch.max(action_label, 1)[1]) # accuracy.append( # int(action_label.argmax(axis=1).eq(action_out.argmax(axis=1)).sum()) / len( # action_label)) # action_loss_val.append(action_l.item()) # # score_l = score_criterion(score_out, score_label.reshape(-1, 1)) # score_loss_val.append(score_l.item()) # print("Vid. {}".format(curr_domain)) # print("\tAccuracy: {}".format(np.mean(accuracy))) # print("\tScore loss: {}".format(np.mean(score_loss_val))) # print("\tAction loss: {}".format(np.mean(action_loss_val))) # if args.visualize: # writer.add_scalars('data/val_video_{}'.format(curr_domain_temp), # {'action_loss_val': np.mean(action_loss_val), # 'score_loss_val': np.mean(score_loss_val), # 'total_val': np.mean(score_loss_val) + np.mean( # action_loss_val), # 'accuracy': np.mean(accuracy)}, # global_step=epoch) if args.visualize: writer.add_scalars('data/epoch_loss', { 'action_loss_tr': action_loss_tr / epoch_size_pos, 'score_loss_tr': score_loss_tr / epoch_size, 'total_tr': action_loss_tr / epoch_size_pos + score_loss_tr / epoch_size }, global_step=epoch) # reset epoch loss counters action_loss_tr = 0 score_loss_tr = 0 # if new epoch (including the first iteration), initialize the batch iterator # or just resuming where batch_iterator_pos and neg haven't been initialized if len(batch_iterators_pos_train) == 0 or len( batch_iterators_neg_train) == 0: # create batch iterator for data_loader_pos in data_loaders_pos_train: batch_iterators_pos_train.append(iter(data_loader_pos)) for data_loader_neg in data_loaders_neg_train: batch_iterators_neg_train.append(iter(data_loader_neg)) # if not batch_iterators_pos_train[curr_domain]: # # create batch iterator # batch_iterators_pos_train[curr_domain] = iter(data_loaders_pos_train[curr_domain]) # # if not batch_iterators_neg_train[curr_domain]: # # create batch iterator # batch_iterators_neg_train[curr_domain] = iter(data_loaders_neg_train[curr_domain]) # load train data if which_dataset[iteration % len(which_dataset)]: # if positive try: images_list, bbox_list, action_labels, score_label, vid_idx = next( batch_iterators_pos_train[curr_domain]) except StopIteration: batch_iterators_pos_train[curr_domain] = iter( data_loaders_pos_train[curr_domain]) images_list, bbox_list, action_labels, score_label, vid_idx = next( batch_iterators_pos_train[curr_domain]) else: try: images_list, bbox_list, action_labels, score_label, vid_idx = next( batch_iterators_neg_train[curr_domain]) except StopIteration: batch_iterators_neg_train[curr_domain] = iter( data_loaders_neg_train[curr_domain]) images_list, bbox_list, action_labels, score_label, vid_idx = next( batch_iterators_neg_train[curr_domain]) # TODO: make sure different obj are paired differenlty, so not always pos with pos if args.cuda: images_list = images_list.to('cuda', non_blocking=True) # bbox = torch.Tensor(bbox.cuda()) action_labels = action_labels.to('cuda', non_blocking=True) score_label = score_label.float().to('cuda', non_blocking=True) else: images = torch.Tensor(images) bbox = torch.Tensor(bbox) action_label = torch.Tensor(action_label) score_label = torch.Tensor(score_label) # TRAIN net.train() action_out, score_out = net(images_list) # load ADNetDomainSpecific with video index if args.cuda: net.module.load_domain_specific(domain_specific_nets[curr_domain]) else: net.load_domain_specific(domain_specific_nets[curr_domain]) # backprop optimizer.zero_grad() accuracy_arr = [] score_l = score_criterion( score_out, torch.cat((score_label.reshape(-1, 1), score_label.reshape(-1, 1)), dim=1)) if which_dataset[iteration % len(which_dataset)]: # if positive action_l = action_criterion( action_out, action_labels.reshape(-1, num_obj_to_track * opts['num_actions'])) loss = action_l + score_l for i in range(num_obj_to_track): accuracy_arr.append(int(action_labels[:, i, :].argmax(axis=1).eq( action_out[:, i * opts['num_actions']:(i + 1) * opts['num_actions']].argmax(axis=1)).sum()) \ / len(action_labels)) else: action_l = -1 accuracy_arr = [-1] * num_obj_to_track loss = score_l loss.backward() optimizer.step() if action_l != -1: action_loss_tr += action_l.item() score_loss_tr += score_l.item() # save the ADNetDomainSpecific back to their module if args.cuda: domain_specific_nets[curr_domain].load_weights_from_adnet( net.module) else: domain_specific_nets[curr_domain].load_weights_from_adnet(net) if args.visualize: if action_l != -1: writer.add_scalars( 'data/iter_loss', { 'action_loss_tr': action_l.item(), 'score_loss_tr': score_l.item(), 'total_tr': (action_l.item() + score_l.item()) }, global_step=iteration) else: writer.add_scalars('data/iter_loss', { 'score_loss_tr': score_l.item(), 'total_tr': score_l.item() }, global_step=iteration) for i in range(num_obj_to_track): accuracy = accuracy_arr[i] if accuracy >= 0: writer.add_scalars('data/iter_acc_{}'.format(i), {'accuracy_tr': accuracy}, global_step=iteration) t1 = time.time() time_arr[iteration % 10] = t1 - t0 if iteration % 10 == 0: # print('Avg. 10 iter time: %.4f sec.' % time_arr.sum()) # print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data.item()), end=' ') if args.visualize and args.send_images_to_visualization: random_batch_index = np.random.randint(images.size(0)) writer.add_image('image', images.data[random_batch_index].cpu().numpy(), random_batch_index) if args.visualize: writer.add_scalars('data/time', {'time_10_it': time_arr.sum()}, global_step=iteration) if iteration % 5000 == 0: print('Saving state, iter:', iteration) domain_specific_nets_state_dict = [] for domain_specific_net in domain_specific_nets: domain_specific_nets_state_dict.append( domain_specific_net.state_dict()) torch.save( { 'epoch': epoch, 'adnet_state_dict': net.state_dict(), 'adnet_domain_specific_state_dict': domain_specific_nets, 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(args.save_folder, args.save_file) + repr(iteration) + '_epoch' + repr(epoch) + '.pth') # final save torch.save( { 'epoch': epoch, 'adnet_state_dict': net.state_dict(), 'adnet_domain_specific_state_dict': domain_specific_nets, 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(args.save_folder, args.save_file) + '.pth') return net, domain_specific_nets, train_videos
def adnet_test_sl(args, opts, mot): if torch.cuda.is_available(): if args.cuda: torch.set_default_tensor_type('torch.cuda.FloatTensor') if not args.cuda: print( "WARNING: It looks like you have a CUDA device, but aren't " + "using CUDA.\nRun with --cuda for optimal training speed.") torch.set_default_tensor_type('torch.FloatTensor') else: torch.set_default_tensor_type('torch.FloatTensor') if not os.path.exists(args.save_folder): os.mkdir(args.save_folder) if args.visualize: writer = SummaryWriter( log_dir=os.path.join('tensorboardx_log', args.save_file)) train_videos = get_train_videos(opts) opts['num_videos'] = len(train_videos['video_names']) net, domain_specific_nets = adnet(opts=opts, trained_file=args.resume, multidomain=args.multidomain) if args.cuda: net = nn.DataParallel(net) cudnn.benchmark = True net = net.cuda() net.eval() action_criterion = nn.CrossEntropyLoss() score_criterion = nn.BCELoss() print('generating Supervised Learning dataset..') # dataset = SLDataset(train_videos, opts, transform= if mot: datasets_pos, datasets_neg = initialize_pos_neg_dataset_mot( train_videos, opts, transform=ADNet_Augmentation(opts)) else: datasets_pos, datasets_neg = initialize_pos_neg_dataset( train_videos, opts, transform=ADNet_Augmentation(opts)) number_domain = opts['num_videos'] assert number_domain == len( datasets_pos ), "Num videos given in opts is incorrect! It should be {}".format( len(datasets_neg)) batch_iterators_pos_val = [] batch_iterators_neg_val = [] # calculating number of data len_dataset_pos = 0 len_dataset_neg = 0 for dataset_pos in datasets_pos: len_dataset_pos += len(dataset_pos) for dataset_neg in datasets_neg: len_dataset_neg += len(dataset_neg) epoch_size_pos = len_dataset_pos // opts['minibatch_size'] epoch_size_neg = len_dataset_neg // opts['minibatch_size'] epoch_size = epoch_size_pos + epoch_size_neg # 1 epoch, how many iterations print("1 epoch = " + str(epoch_size) + " iterations") max_iter = opts['numEpoch'] * epoch_size print("maximum iteration = " + str(max_iter)) data_loaders_pos_val = [] data_loaders_neg_val = [] for dataset_pos in datasets_pos: data_loaders_pos_val.append( data.DataLoader(dataset_pos, opts['minibatch_size'], num_workers=2, shuffle=True, pin_memory=True)) for dataset_neg in datasets_neg: data_loaders_neg_val.append( data.DataLoader(dataset_neg, opts['minibatch_size'], num_workers=2, shuffle=True, pin_memory=True)) net.eval() for curr_domain in range(number_domain): accuracy = [] action_loss_val = [] score_loss_val = [] # load ADNetDomainSpecific with video index if args.cuda: net.module.load_domain_specific(domain_specific_nets[curr_domain]) else: net.load_domain_specific(domain_specific_nets[curr_domain]) for i, temp in enumerate([ data_loaders_pos_val[curr_domain], data_loaders_neg_val[curr_domain] ]): dont_show = False for images, bbox, action_label, score_label, indices in tqdm(temp): images = images.to('cuda', non_blocking=True) action_label = action_label.to('cuda', non_blocking=True) score_label = score_label.float().to('cuda', non_blocking=True) # forward action_out, score_out = net(images) if i == 0: # if positive action_l = action_criterion(action_out, torch.max(action_label, 1)[1]) action_loss_val.append(action_l.item()) accuracy.append( int( action_label.argmax(axis=1).eq( action_out.argmax(axis=1)).sum()) / len(action_label)) score_l = score_criterion(score_out, score_label.reshape(-1, 1)) score_loss_val.append(score_l.item()) if args.display_images and not dont_show: if i == 0: dataset = datasets_pos[curr_domain] color = (0, 255, 0) conf = 1 else: dataset = datasets_neg[curr_domain] color = (0, 0, 255) conf = 0 for j, index in enumerate(indices): im = cv2.imread(dataset.train_db['img_path'][index]) bbox = dataset.train_db['bboxes'][index] action_label = np.array( dataset.train_db['labels'][index]) cv2.rectangle(im, (bbox[0], bbox[1]), (bbox[0] + bbox[2], bbox[1] + bbox[3]), color, 2) print("\n\nTarget actions: {}".format( action_label.argmax())) print("Predicted actions: {}".format( action_out.data[j].argmax())) print("Target conf: {}".format(conf)) print("Predicted conf: {}".format(score_out.data[j])) # print("Score loss: {}".format(score_l.item())) # print("Action loss: {}".format(action_l.item())) cv2.imshow("Test", im) key = cv2.waitKey(0) & 0xFF # if the `q` key was pressed, break from the loop if key == ord("q"): dont_show = True break elif key == ord("s"): cv2.imwrite( "vid {} t:{} p:{} c:{}.png".format( curr_domain, action_label.argmax(), action_out.data[i].argmax(), score_out.data[i].item()), im) print("Vid. {}".format(curr_domain)) print("\tAccuracy: {}".format(np.mean(accuracy))) print("\tScore loss: {}".format(np.mean(score_loss_val))) print("\tAction loss: {}".format(np.mean(action_loss_val))) sys.exit(0) return net, domain_specific_nets, train_videos