# print('\n==> Freeze backbone...') # for param in model.parameters(): # param.requires_grad = False # criterion = nn.CrossEntropyLoss().cuda() if args.no_mean_norm and not args.std_norm: norm_method = Normalize([0, 0, 0], [1, 1, 1]) elif not args.std_norm: norm_method = Normalize(args.mean, [1, 1, 1]) else: norm_method = Normalize(args.mean, args.std) assert args.train_crop in ['random', 'corner', 'center'] if args.train_crop == 'random': crop_method = MultiScaleRandomCrop(args.scales, args.sample_size) elif args.train_crop == 'corner': crop_method = MultiScaleCornerCrop(args.scales, args.sample_size) elif args.train_crop == 'center': crop_method = MultiScaleCornerCrop(args.scales, args.sample_size, crop_positions=['c']) print('==> Preparing dataset...') print('\nTraining set') # spatial_transform = Compose([ # Gaussian_blur(p=0.5), # RandomHorizontalFlip(), # crop_method, # ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.25, p=0.3), # ToTensor(args.norm_value), norm_method
# print(model) criterion = nn.CrossEntropyLoss() if not opt.no_cuda: criterion = criterion.cuda() if opt.no_mean_norm and not opt.std_norm: norm_method = Normalize([0, 0, 0], [1, 1, 1]) elif not opt.std_norm: norm_method = Normalize(opt.mean, [1, 1, 1]) else: norm_method = Normalize(opt.mean, opt.std) if not opt.no_train: assert opt.train_crop in ['random', 'corner', 'center'] if opt.train_crop == 'random': crop_method = MultiScaleRandomCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'corner': crop_method = MultiScaleCornerCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'center': crop_method = MultiScaleCornerCrop(opt.scales, opt.sample_size, crop_positions=['c']) spatial_transform = Compose([ crop_method, RandomHorizontalFlip(), ToTensor(opt.norm_value), norm_method ]) temporal_transform = TemporalRandomCrop(opt.sample_duration) target_transform = ClassLabel() training_data = get_training_set(opt, spatial_transform, temporal_transform, target_transform)
criterion = criterion.cuda() if opt.no_mean_norm and not opt.std_norm: norm_method = Normalize([0, 0, 0], [1, 1, 1]) elif not opt.std_norm: norm_method = Normalize(opt.mean, [1, 1, 1]) else: norm_method = Normalize(opt.mean, opt.std) if not opt.no_train: ##-------------------------------------------------------------------------------------------- if opt.model == 'I3D': assert opt.train_crop in ['random', 'corner', 'center'] if opt.train_crop == 'random': crop_method = MultiScaleRandomCrop([0.875], opt.sample_size) elif opt.train_crop == 'corner': crop_method = MultiScaleCornerCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'center': crop_method = MultiScaleCornerCrop(opt.scales, opt.sample_size, crop_positions=['c']) spatial_transform = Compose([ Scale((256, 256)), crop_method, RandomHorizontalFlip(), ToTensor(opt.norm_value), norm_method ]) temporal_transform = TemporalRandomCrop(opt.sample_duration, 1) elif opt.model == 'resnet_50': assert opt.train_crop in ['random', 'corner', 'center'] if opt.train_crop == 'random':
def main(): global epochs # Config parser = argparse.ArgumentParser( description="To read EgoGesture Dataset and run through SSAR network") parser.add_argument('--path', default='', help='full path to EgoGesture Dataset') args = parser.parse_args() path = args.path # Setup multiscale random crop scales = [initial_scale] for _ in range(1, n_scales): scales.append(scales[-1] * scale_step) # Setup datasets / dataloaders if do_data_augmentation: train_spatial_transforms = Compose([ MultiScaleRandomCrop(scales, (126, 224)), SpatialElasticDisplacement() ]) else: train_spatial_transforms = transforms.Resize((126, 224)) image_transform_train = Compose([ train_spatial_transforms, transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) image_transform_val = Compose([ transforms.Resize((126, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) image_transform_test = Compose([ transforms.Resize((126, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) mask_transform = Compose([train_spatial_transforms, transforms.ToTensor()]) if not do_data_augmentation: image_transform_train = image_transform_val subject_ids_train = [ 3, 4, 5, 6, 8, 10, 15, 16, 17, 20, 21, 22, 23, 25, 26, 27, 30, 32, 36, 38, 39, 40, 42, 43, 44, 45, 46, 48, 49, 50 ] subject_ids_val = [1, 7, 12, 13, 24, 29, 33, 34, 35, 37] subject_ids_test = [2, 9, 11, 14, 18, 19, 28, 31, 41, 47] if mode == 'training': train_dataset = EgoGestDataSequence(path, 'train_dataset', image_transform_train, mask_transform, get_mask=use_mask_loss, subject_ids=subject_ids_train) val_dataset = EgoGestDataSequence(path, 'val_dataset', image_transform_val, mask_transform, get_mask=use_mask_loss, subject_ids=subject_ids_val) # If we're not in training mode then switch the training dataset out with test or validation elif mode == 'validation': train_dataset = EgoGestDataSequence(path, 'val_dataset', image_transform_val, mask_transform, get_mask=use_mask_loss, subject_ids=subject_ids_val) else: train_dataset = EgoGestDataSequence(path, 'val_dataset', image_transform_test, mask_transform, get_mask=use_mask_loss, subject_ids=subject_ids_test) # train_indices, val_indices, test_indices = check_and_split_data(host_name=hostname, # data_folder=path, # dataset_len=len(dataset), # train_fraction=0.6, # validation_fraction=0.2) torch.manual_seed(42) torch.backends.cudnn.deterministic = True train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=True, shuffle=True, collate_fn=collate_fn_padd) if mode == 'training': val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=True, shuffle=True, collate_fn=collate_fn_padd) # Init model and load pre-trained weights rnet = resnet.resnet18(False) model = SSAR(ResNet=rnet, input_size=83, number_of_classes=83, batch_size=batch_size, dropout=dropout).cuda() model_weights = './weights/final_weights.pth' state = model.state_dict() loaded_weights = torch.load(model_weights) state.update(loaded_weights) model.load_state_dict(state) # Setup optimizer and loss criterion = torch.nn.CrossEntropyLoss(ignore_index=label_mask_value) criterion = criterion.cuda() if mode == 'training': set_train_mode( model, train=True ) # Need this here so the optimizer has the correct parameters to be trained optimizer = Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=learning_rate) else: optimizer = None # Continue from previous training checkpoint epoch_resume, step_resume, best_val_loss = load_latest( model, results_path, training_mode, optimizer) if not restore_training_variables: epoch_resume = 0 step_resume = 0 best_val_loss = np.inf # Train / test / val setup if mode != 'training': epoch_resume = 0 step_resume = 0 epochs = 1 # old_tensor_set = set() # Accuracy bar plot plt.ion() fig, ax = plt.subplots(nrows=3, ncols=1, figsize=(10, 7), dpi=100) train_acc_bars = ax[0].bar(rel_poses, 0, 1 / accuracy_bins) val_acc_bars = ax[1].bar(rel_poses, 0, 1 / accuracy_bins) train_loss_line, = ax[2].plot([], [], label='Train Loss') val_loss_line, = ax[2].plot([], [], label='Val Loss') train_acc_texts = [ ax[0].text(x, y, "", horizontalalignment='center', verticalalignment='bottom') for x, y in zip(rel_poses, np.ones_like(rel_poses)) ] val_acc_texts = [ ax[1].text(x, y, "", horizontalalignment='center', verticalalignment='bottom') for x, y in zip(rel_poses, np.ones_like(rel_poses)) ] ax[0].set_ylim(0., 1.1) ax[0].set_title('Relative Gesture Position vs Training Accuracy') ax[1].set_ylim(0., 1.1) ax[1].set_title('Relative Gesture Position vs Validation Accuracy') ax[2].legend(loc='best') # loss_text = ax[2].text(0, -0.2, "Loss: ") plt.show() # Setup movie moviewriter for writing accuracy plot over time # moviewriter = FFMpegWriter(fps=1) # moviewriter.setup(fig, os.path.join(results_path, 'accuracy_over_time.mp4'), dpi=100) # Main training loop if mode == 'training': optimizer.zero_grad() train_history = {} val_history = {} patience_counter = 0 for epoch in range(epoch_resume, epochs): # Display info print(f"Epoch: {epoch}") if epoch == epoch_resume and step_resume > 0: print(f"Fast forwarding to train step {step_resume}") # Reset epoch stats train_metrics = {} # Switch to train mode while freezing parts of the model we don't want to train set_train_mode(model, train=True) # Train if mode == 'training': print('Training:') for train_step, batch in enumerate(train_loader): # Advance train_loader to resume training from last checkpointed position (Note: Assumes same batch size) if epoch == epoch_resume and train_step < step_resume: del batch continue # Save model if mode == 'training' and train_step % 100 == 0 and ( train_step != step_resume or epoch != epoch_resume): save_model(model, optimizer, training_mode, epoch, train_step, best_val_loss, results_path) # Do one training train_step (may not actually train_step optimizer if doing gradiant accumulation) loss, batch_correct_count_samples = process_batch(model, train_step, batch, criterion, optimizer, mode=mode) del batch # Update metrics update_metrics(train_metrics, epoch, loss, batch_correct_count_samples) if (train_step + 1) % 10 == 0: # Display metrics print_metrics(train_metrics, train_step) update_accuracy_plot(train_acc_bars, train_acc_texts, train_metrics['accuracy_hist']) # Update train metric history and plots for this epoch update_epoch_history(train_history, train_metrics) update_loss_plot(train_loss_line, train_history) # Validation if mode == 'training': print('Validation:') val_metrics = {} # Switch to evaluation mode for validation set_train_mode(model, train=False) for val_step, batch in enumerate(val_loader): loss, batch_correct_count_samples = process_batch( model, val_step, batch, criterion, optimizer, mode='validation') # Update metrics update_metrics(val_metrics, epoch, loss, batch_correct_count_samples) if (val_step + 1) % 10 == 0: # Display metrics print_metrics(val_metrics, val_step) update_accuracy_plot(val_acc_bars, val_acc_texts, val_metrics['accuracy_hist']) # Update validation metric history and plots update_epoch_history(val_history, val_metrics) update_loss_plot(val_loss_line, val_history) # Early stoping if val_metrics['loss_epoch'] < best_val_loss: best_val_loss = val_metrics['loss_epoch'] patience_counter = 0 save_model(model, optimizer, training_mode, epoch, val_step, best_val_loss, results_path, filename_override='model_best.pth') else: patience_counter += 1 if patience_counter >= early_stoppping_patience: print( f'Validation accuracy did not improve for {patience_counter} epochs, stopping' ) break # Save final model if mode == 'training' and (train_step != step_resume or epoch != epoch_resume): save_model(model, optimizer, training_mode, epoch, train_step, best_val_loss, results_path) print('Done!') plt.ioff() plt.show()
scheduler = optim.lr_scheduler.StepLR(optimizer, 1, gamma=decay_rate) criterion = nn.CrossEntropyLoss() if not opt.no_cuda: criterion = criterion.cuda() if opt.no_mean_norm and not opt.std_norm: norm_method = Normalize([0, 0, 0], [1, 1, 1]) elif not opt.std_norm: norm_method = Normalize(opt.mean, [1, 1, 1]) else: norm_method = Normalize(opt.mean, opt.std) assert opt.train_crop in ['random', 'corner', 'center'] if opt.train_crop == 'random': ucf_crop = MultiScaleRandomCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'corner': ucf_crop = MultiScaleCornerCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'center': ucf_crop = MultiScaleCornerCrop(opt.scales, opt.sample_size, crop_positions=['c']) spatial_transform = [] temporal_transform = [] target_transform = [] ucf_transform = Compose([ RandomHorizontalFlip(), ucf_crop, ToTensor(opt.norm_value), norm_method, ])
parameters, lr=opt.learning_rate, momentum=opt.momentum, dampening=dampening, weight_decay=opt.weight_decay, nesterov=opt.nesterov) normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ## prepare train if not opt.no_train: temporal_transform = TemporalSegmentRandomCrop(opt.segment_number, opt.sample_duration) assert opt.train_crop in ['random', 'corner', 'center'] if opt.train_crop == 'random': spatial_crop_method = MultiScaleRandomCrop(opt.scales, opt.frame_size) elif opt.train_crop == 'corner': spatial_crop_method = MultiScaleCornerCrop(opt.scales, opt.frame_size) elif opt.train_crop == 'center': spatial_crop_method = MultiScaleCornerCrop(opt.scales, opt.frame_size, crop_positions=['c']) spatial_transform = Compose([ spatial_crop_method, RandomHorizontalFlip(), ToTensor(opt.norm_value), normalize ]) training_data = get_training_set(opt, spatial_transform, temporal_transform) train_loader = DataLoaderX( training_data, batch_size=opt.batch_size,
elif opt.optimizer == 'radam': optimizer = RAdam( parameters, lr=opt.learning_rate, betas=(0.9, 0.999), weight_decay=opt.weight_decay) normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ## prepare train if not opt.no_train: temporal_transform = TemporalSegmentRandomCrop(opt.segment_number, opt.sample_duration) assert opt.train_crop in ['random', 'corner', 'center'] if opt.train_crop == 'random': sceobj_crop_method = MultiScaleRandomCrop(opt.scales, opt.sceobj_frame_size) elif opt.train_crop == 'corner': sceobj_crop_method = MultiScaleCornerCrop(opt.scales, opt.sceobj_frame_size) elif opt.train_crop == 'center': sceobj_crop_method = MultiScaleCornerCrop(opt.scales, opt.sceobj_frame_size, crop_positions=['c']) sceobj_spatial_transform = Compose([ sceobj_crop_method, RandomHorizontalFlip(), ToTensor(opt.norm_value), normalize ]) #sceobj_spatial_transform = transforms.Compose([ # transforms.Resize(256), # transforms.CenterCrop(opt.sceobj_frame_size), # transforms.RandomHorizontalFlip(), # transforms.ToTensor(),
def create_dataloader(args): if args.root_path != '': args.video_path = os.path.join(args.root_path, args.video_path) args.annotation_path = os.path.join(args.root_path, args.annotation_path) args.result_path = os.path.join(args.root_path, args.result_path) if args.resume_path: args.resume_path = os.path.join(args.root_path, args.resume_path) if args.pretrain_path: # args.pretrain_path = os.path.join(args.root_path, args.pretrain_path) args.pretrain_path = os.path.abspath(args.pretrain_path) args.scales = [args.initial_scale] for i in range(1, args.n_scales): args.scales.append(args.scales[-1] * args.scale_step) args.mean = get_mean(args.norm_value, dataset=args.mean_dataset) args.std = get_std(args.norm_value) if args.no_mean_norm and not args.std_norm: norm_method = Normalize([0, 0, 0], [1, 1, 1]) elif not args.std_norm: norm_method = Normalize(args.mean, [1, 1, 1]) else: norm_method = Normalize(args.mean, args.std) assert args.train_crop in ['random', 'corner', 'center'] if args.train_crop == 'random': crop_method = MultiScaleRandomCrop(args.scales, args.sample_size) elif args.train_crop == 'corner': crop_method = MultiScaleCornerCrop(args.scales, args.sample_size) elif args.train_crop == 'center': crop_method = MultiScaleCornerCrop(args.scales, args.sample_size, crop_positions=['c']) spatial_transform = Compose([ crop_method, RandomHorizontalFlip(), ToTensor(args.norm_value), norm_method ]) temporal_transform = TemporalRandomCrop(args.sample_duration) target_transform = ClassLabel() training_data = get_training_set(args, spatial_transform, temporal_transform, target_transform) train_loader = torch.utils.data.DataLoader(training_data, batch_size=args.batch_size, shuffle=True, num_workers=args.n_threads, pin_memory=True) spatial_transform = Compose([ # Scale(args.sample_size), Scale(int(args.sample_size / args.scale_in_test)), # CenterCrop(args.sample_size), CornerCrop(args.sample_size, args.crop_position_in_test), ToTensor(args.norm_value), norm_method ]) temporal_transform = TemporalCenterCrop(args.sample_duration) target_transform = ClassLabel() validation_data = get_validation_set(args, spatial_transform, temporal_transform, target_transform) val_loader = torch.utils.data.DataLoader(validation_data, batch_size=1, shuffle=False, num_workers=args.n_threads, pin_memory=True) return train_loader, val_loader
def main(): resnet_in = generate_model(opt) resnet_in.module.fc = Identity() model = ReNet34(resnet_in, encode_length=encode_length) if opt.no_mean_norm and not opt.std_norm: norm_method = Normalize([0, 0, 0], [1, 1, 1]) elif not opt.std_norm: norm_method = Normalize(opt.mean, [1, 1, 1]) else: norm_method = Normalize(opt.mean, opt.std) if not opt.no_train: assert opt.train_crop in ['random', 'corner', 'center'] if opt.train_crop == 'random': crop_method = MultiScaleRandomCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'corner': crop_method = MultiScaleCornerCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'center': crop_method = MultiScaleCornerCrop(opt.scales, opt.sample_size, crop_positions=['c']) ## train loader spatial_transform = Compose([ crop_method, RandomHorizontalFlip(), ToTensor(opt.norm_value), norm_method ]) temporal_transform = TemporalRandomCrop(opt.sample_duration) target_transform = ClassLabel() training_data = get_training_set(opt, spatial_transform, temporal_transform, target_transform) train_loader = torch.utils.data.DataLoader(training_data, batch_size=opt.batch_size, shuffle=True, num_workers=opt.n_threads, pin_memory=True) ## test loader spatial_transform = Compose([ Scale(int(opt.sample_size / opt.scale_in_test)), CornerCrop(opt.sample_size, opt.crop_position_in_test), ToTensor(opt.norm_value), norm_method ]) temporal_transform = LoopPadding(opt.sample_duration) target_transform = ClassLabel() test_data = get_test_set(opt, spatial_transform, temporal_transform, target_transform) test_loader = torch.utils.data.DataLoader(test_data, batch_size=opt.batch_size, shuffle=False, num_workers=opt.n_threads, pin_memory=True) ## Database loader spatial_transform = Compose([ Scale(int(opt.sample_size / opt.scale_in_test)), CornerCrop(opt.sample_size, opt.crop_position_in_test), ToTensor(opt.norm_value), norm_method ]) temporal_transform = LoopPadding(opt.sample_duration) target_transform = ClassLabel() validation_data = get_validation_set(opt, spatial_transform, temporal_transform, target_transform) database_loader = torch.utils.data.DataLoader( validation_data, batch_size=opt.batch_size, shuffle=False, num_workers=opt.n_threads, pin_memory=True) if opt.nesterov: dampening = 0 else: dampening = opt.dampening optimizer = optim.SGD(model.parameters(), lr=opt.learning_rate, momentum=opt.momentum, dampening=dampening, weight_decay=opt.weight_decay, nesterov=opt.nesterov) scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=opt.lr_patience) if opt.resume_path: print('loading checkpoint {}'.format(opt.resume_path)) checkpoint = torch.load(opt.resume_path) assert opt.arch == checkpoint['arch'] opt.begin_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) if not opt.no_train: optimizer.load_state_dict(checkpoint['optimizer']) for state in optimizer.state.values(): for k, v in state.items(): if torch.is_tensor(v): state[k] = v.cuda() print('run') for epoch in range(opt.begin_epoch, opt.n_epochs + 1): model.cuda().train() for i, (images, labels) in enumerate(train_loader): images = Variable(images.cuda()) labels = Variable(labels.cuda().long()) # Forward + Backward + Optimize optimizer.zero_grad() x, _, b = model(images) target_b = F.cosine_similarity(b[:int(labels.size(0) / 2)], b[int(labels.size(0) / 2):]) target_x = F.cosine_similarity(x[:int(labels.size(0) / 2)], x[int(labels.size(0) / 2):]) loss = F.mse_loss(target_b, target_x) loss.backward() optimizer.step() scheduler.step() # Test the Model if (epoch + 1) % 10 == 0: model.eval() retrievalB, retrievalL, queryB, queryL = compress( database_loader, test_loader, model) result_map = calculate_top_map(qB=queryB, rB=retrievalB, queryL=queryL, retrievalL=retrievalL, topk=100) print('--------mAP@100: {}--------'.format(result_map))
def objective(trial): opt = parse_opts() if trial: opt.weight_decay = trial.suggest_uniform('weight_decay', 0.01, 0.1) opt.learning_rate = trial.suggest_uniform('learning_rate', 1 - 5, 1 - 4) if opt.root_path != '': opt.video_path = os.path.join(opt.root_path, opt.video_path) opt.annotation_path = os.path.join(opt.root_path, opt.annotation_path) opt.result_path = os.path.join(opt.root_path, opt.result_path) if opt.resume_path: opt.resume_path = os.path.join(opt.root_path, opt.resume_path) if opt.pretrain_path: opt.pretrain_path = os.path.join(opt.root_path, opt.pretrain_path) opt.scales = [opt.initial_scale] for i in range(1, opt.n_scales): opt.scales.append(opt.scales[-1] * opt.scale_step) opt.arch = '{}-{}'.format(opt.model, opt.model_depth) opt.mean = get_mean(opt.norm_value, dataset=opt.mean_dataset) opt.std = get_std(opt.norm_value) print(opt) with open(os.path.join(opt.result_path, 'opts.json'), 'w') as opt_file: json.dump(vars(opt), opt_file) torch.manual_seed(opt.manual_seed) model, parameters = generate_model(opt) print(model) criterion = nn.CrossEntropyLoss() if not opt.no_cuda: criterion = criterion.cuda() if opt.no_mean_norm and not opt.std_norm: norm_method = Normalize([0, 0, 0], [1, 1, 1]) elif not opt.std_norm: norm_method = Normalize(opt.mean, [1, 1, 1]) else: norm_method = Normalize(opt.mean, opt.std) # norm_method = Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) if not opt.no_train: assert opt.train_crop in ['random', 'corner', 'center'] if opt.train_crop == 'random': crop_method = MultiScaleRandomCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'corner': crop_method = MultiScaleCornerCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'center': crop_method = MultiScaleCornerCrop(opt.scales, opt.sample_size, crop_positions=['c']) spatial_transform = Compose([ crop_method, RandomHorizontalFlip(), ToTensor(opt.norm_value), norm_method ]) temporal_transform = TemporalRandomCrop(opt.sample_duration) target_transform = ClassLabel() training_data = get_training_set(opt, spatial_transform, temporal_transform, target_transform) train_loader = torch.utils.data.DataLoader( training_data, batch_size=opt.batch_size, # sampler option is mutually exclusive with shuffle shuffle=False, sampler=ImbalancedDatasetSampler(training_data), num_workers=opt.n_threads, pin_memory=True) train_logger = Logger(os.path.join(opt.result_path, 'train.log'), ['epoch', 'loss', 'acc', 'lr']) train_batch_logger = Logger( os.path.join(opt.result_path, 'train_batch.log'), ['epoch', 'batch', 'iter', 'loss', 'acc', 'lr']) optimizer = optim.Adam(parameters, lr=opt.learning_rate, weight_decay=opt.weight_decay) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, verbose=True, factor=0.1**0.5) if not opt.no_val: spatial_transform = Compose([ Scale(opt.sample_size), CenterCrop(opt.sample_size), ToTensor(opt.norm_value), norm_method ]) temporal_transform = LoopPadding(opt.sample_duration) target_transform = ClassLabel() validation_data = get_validation_set(opt, spatial_transform, temporal_transform, target_transform) val_loader = torch.utils.data.DataLoader( validation_data, batch_size=opt.batch_size, shuffle=False, sampler=ImbalancedDatasetSampler(validation_data), num_workers=opt.n_threads, pin_memory=True) val_logger = Logger(os.path.join(opt.result_path, 'val.log'), ['epoch', 'loss', 'acc']) if opt.resume_path: print('loading checkpoint {}'.format(opt.resume_path)) checkpoint = torch.load(opt.resume_path) assert opt.arch == checkpoint['arch'] opt.begin_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) if not opt.no_train: optimizer.load_state_dict(checkpoint['optimizer']) print('run') writer = SummaryWriter( comment= f"_wd{opt.weight_decay}_lr{opt.learning_rate}_ft_begin{opt.ft_begin_index}_pretrain{not opt.pretrain_path == ''}" ) for i in range(opt.begin_epoch, opt.n_epochs + 1): if not opt.no_train: epoch, losses_avg, accuracies_avg = train_epoch( i, train_loader, model, criterion, optimizer, opt, train_logger, train_batch_logger) writer.add_scalar('loss/train', losses_avg, epoch) writer.add_scalar('acc/train', accuracies_avg, epoch) if not opt.no_val: epoch, val_losses_avg, val_accuracies_avg = val_epoch( i, val_loader, model, criterion, opt, val_logger) writer.add_scalar('loss/val', val_losses_avg, epoch) writer.add_scalar('acc/val', val_accuracies_avg, epoch) if not opt.no_train and not opt.no_val: scheduler.step(val_losses_avg) print('=' * 100) if opt.test: spatial_transform = Compose([ Scale(int(opt.sample_size / opt.scale_in_test)), CornerCrop(opt.sample_size, opt.crop_position_in_test), ToTensor(opt.norm_value), norm_method ]) temporal_transform = LoopPadding(opt.sample_duration) target_transform = VideoID() test_data = get_test_set(opt, spatial_transform, temporal_transform, target_transform) test_loader = torch.utils.data.DataLoader(test_data, batch_size=opt.batch_size, shuffle=False, num_workers=opt.n_threads, pin_memory=True) test.test(test_loader, model, opt, test_data.class_names) writer.close() return val_losses_avg
labels = np.load(args.pesudo_label_file) # get teacher and student model student_model = create_model(args) teacher_model = create_model(args, ema=True) """ Prepare dataset and loader for mean-teacher training Here each batch contains three parts, each has same number of samples. Parts 1: ucf101 samples, contribute to cross_entropy loss Parts 2&3: kinetics samples with different augmentation, contribute to consistency loss right now only consider augmentation on temporal dimension, since i think temporal information contributes more to action recognition than spatial information """ norm_method = Normalize(args.mean, [1, 1, 1]) ucf_crop = MultiScaleRandomCrop(args.scales, args.sample_size) ucf_spatial = Compose([ RandomHorizontalFlip(), ucf_crop, ToTensor(args.norm_value), norm_method, ]) kinetics_spatial = Compose([ RandomHorizontalFlip(), ucf_crop, ToTensor(args.norm_value), norm_method, ]) ucf_temporal = TemporalRandomCrop(args.sample_duration, args.downsample) # ucf_temporal = TemporalCenterCrop(args.sample_duration, args.downsample)