def test_pksampler(self): p, k = 16, 4 # Ensure sampler does not allow p to be greater than num_classes dataset = FakeData(size=100, num_classes=10, image_size=(3, 1, 1)) targets = [target.item() for _, target in dataset] self.assertRaises(AssertionError, PKSampler, targets, p, k) # Ensure p, k constraints on batch trans = transforms.Compose([ transforms.PILToTensor(), transforms.ConvertImageDtype(torch.float), ]) dataset = FakeData(size=1000, num_classes=100, image_size=(3, 1, 1), transform=trans) targets = [target.item() for _, target in dataset] sampler = PKSampler(targets, p, k) loader = DataLoader(dataset, batch_size=p * k, sampler=sampler) for _, labels in loader: bins = defaultdict(int) for label in labels.tolist(): bins[label] += 1 # Ensure that each batch has samples from exactly p classes self.assertEqual(len(bins), p) # Ensure that there are k samples from each class for b in bins: self.assertEqual(bins[b], k)
def main(args): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") p = args.labels_per_batch k = args.samples_per_label batch_size = p * k model = EmbeddingNet() if args.resume: model.load_state_dict(torch.load(args.resume)) model.to(device) criterion = TripletMarginLoss(margin=args.margin) optimizer = Adam(model.parameters(), lr=args.lr) transform = transforms.Compose([ transforms.Lambda(lambda image: image.convert("RGB")), transforms.Resize((224, 224)), transforms.PILToTensor(), transforms.ConvertImageDtype(torch.float), ]) # Using FMNIST to demonstrate embedding learning using triplet loss. This dataset can # be replaced with any classification dataset. train_dataset = FashionMNIST(args.dataset_dir, train=True, transform=transform, download=True) test_dataset = FashionMNIST(args.dataset_dir, train=False, transform=transform, download=True) # targets is a list where the i_th element corresponds to the label of i_th dataset element. # This is required for PKSampler to randomly sample from exactly p classes. You will need to # construct targets while building your dataset. Some datasets (such as ImageFolder) have a # targets attribute with the same format. targets = train_dataset.targets.tolist() train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=PKSampler(targets, p, k), num_workers=args.workers) test_loader = DataLoader(test_dataset, batch_size=args.eval_batch_size, shuffle=False, num_workers=args.workers) for epoch in range(1, args.epochs + 1): print("Training...") train_epoch(model, optimizer, criterion, train_loader, device, epoch, args.print_freq) print("Evaluating...") evaluate(model, test_loader, device) print("Saving...") save(model, epoch, args.save_dir, "ckpt.pth")
def main(args): device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') p = args.labels_per_batch k = args.samples_per_label batch_size = p * k model = EmbeddingNet(backbone=None, pretrained=args.pretrained) print('pretrained: ', args.pretrained) if args.resume: model.load_state_dict(torch.load(args.resume)) model.to(device) criterion = TripletMarginLoss(margin=args.margin) optimizer = Adam(model.parameters(), lr=args.lr) train_transform = transforms.Compose([transforms.Lambda(lambda image: image.convert('RGB')), transforms.RandomHorizontalFlip(), transforms.Resize((224, 224)), transforms.ToTensor()]) test_transform = transforms.Compose([transforms.Lambda(lambda image: image.convert('RGB')), transforms.Resize((224, 224)), transforms.ToTensor()]) # Using FMNIST to demonstrate embedding learning using triplet loss. This dataset can # be replaced with any classification dataset. train_dataset = ImageFolder(os.path.join(args.dataset_dir, 'train'), transform=train_transform) test_dataset = ImageFolder(os.path.join(args.dataset_dir, 'test'), transform=test_transform) # targets is a list where the i_th element corresponds to the label of i_th dataset element. # This is required for PKSampler to randomly sample from exactly p classes. You will need to # construct targets while building your dataset. Some datasets (such as ImageFolder) have a # targets attribute with the same format. targets = train_dataset.targets train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=PKSampler(targets, p, k), num_workers=args.workers) test_loader = DataLoader(test_dataset, batch_size=args.eval_batch_size, shuffle=False, num_workers=args.workers) for epoch in range(1, args.epochs + 1): print('Training...') train_epoch(model, optimizer, criterion, train_loader, device, epoch, args.print_freq) print('Evaluating...') evaluate(model, test_loader, device) print('Saving...') save(model, epoch, args.save_dir, 'ckpt.pth')
def get_Sampler(sampler, dataset, p=15, k=20): if sampler == 'all': return PKSampler2(dataset, p=p, k=k) else: return PKSampler(dataset, p=p, k=k)
def main(): args = parse_args() conf = Config(args.conf) data_dir = conf.data_dir fold_id = conf.fold_id workspace = Workspace(conf.run_id).setup() workspace.save_conf(args.conf) workspace.log(f'{conf.to_dict()}') torch.cuda.set_device(0) if conf.use_augmentor: if conf.augmentor_type == 'v1': augmentor = create_augmentor_v1( enable_random_morph=conf.enable_random_morph) elif conf.augmentor_type == 'v2': augmentor = create_augmentor_v2( enable_random_morph=conf.enable_random_morph, invert_color=conf.invert_color) elif conf.augmentor_type == 'v3': if conf.input_size_tuple: input_size = tuple(conf.input_size_tuple) else: input_size = (conf.input_size, conf.input_size) if conf.input_size else \ (SOURCE_IMAGE_HEIGHT, SOURCE_IMAGE_WIDTH) augmentor = create_augmentor_v3( input_size, enable_random_morph=conf.enable_random_morph, invert_color=conf.invert_color) else: raise ValueError(conf.augmentor_type) workspace.log(f'Use augmentor: {conf.augmentor_type}') else: augmentor = None if not conf.input_size_tuple and conf.input_size == 0: train_transformer = create_transformer_v1(augmentor=augmentor) val_transformer = create_testing_transformer_v1() workspace.log('Input size: default') else: if conf.input_size_tuple: input_size = tuple(conf.input_size_tuple) else: input_size = (conf.input_size, conf.input_size) train_transformer = create_transformer_v1(input_size=input_size, augmentor=augmentor) val_transformer = create_testing_transformer_v1(input_size=input_size) workspace.log(f'Input size: {input_size}') train_dataset, val_dataset = bengali_dataset( data_dir, fold_id=fold_id, train_transformer=train_transformer, val_transformer=val_transformer, invert_color=conf.invert_color, n_channel=conf.n_channel, use_grapheme_code=conf.use_grapheme_code, logger=workspace.logger) workspace.log(f'#train={len(train_dataset)}, #val={len(val_dataset)}') train_dataset.set_low_freq_groups(n_class=conf.n_class_low_freq) if conf.sampler_type == 'pk': sampler = PKSampler(train_dataset, n_iter_per_epoch=conf.n_iter_per_epoch, p=conf.batch_p, k=conf.batch_k) train_loader = DataLoader(train_dataset, shuffle=False, num_workers=8, pin_memory=True, batch_sampler=sampler) workspace.log(f'{sampler} is enabled') workspace.log(f'Real batch_size={sampler.batch_size}') elif conf.sampler_type == 'random+append': batch_sampler = LowFreqSampleMixinBatchSampler( train_dataset, conf.batch_size, n_low_freq_samples=conf.n_low_freq_samples, drop_last=True) train_loader = DataLoader(train_dataset, shuffle=False, num_workers=8, pin_memory=True, batch_sampler=batch_sampler) workspace.log(f'{batch_sampler} is enabled') workspace.log(f'Real batch_size={batch_sampler.batch_size}') elif conf.sampler_type == 'random': train_loader = DataLoader(train_dataset, batch_size=conf.batch_size, shuffle=True, num_workers=8, pin_memory=True, drop_last=True) else: raise ValueError(f'Invalid sampler_type: {conf.sampler_type}') val_loader = DataLoader(val_dataset, batch_size=conf.batch_size, shuffle=False, num_workers=8, pin_memory=True) workspace.log(f'Create init model: arch={conf.arch}') model = create_init_model(conf.arch, pretrained=True, pooling=conf.pooling_type, dim=conf.feat_dim, use_maxblurpool=conf.use_maxblurpool, remove_last_stride=conf.remove_last_stride, n_channel=conf.n_channel) if conf.weight_file: pretrained_weight = torch.load(conf.weight_file, map_location='cpu') result = model.load_state_dict(pretrained_weight) workspace.log(f'Pretrained weights were loaded: {conf.weight_file}') workspace.log(result) model = model.cuda() sub_models = [] criterion_g = get_criterion(conf.loss_type_g, weight=train_dataset.get_class_weights_g(), rate=conf.ohem_rate) workspace.log(f'Loss type (g): {conf.loss_type_g}') criterion_v = get_criterion(conf.loss_type_v, weights=train_dataset.get_class_weights_v(), rate=conf.ohem_rate) workspace.log(f'Loss type (v): {conf.loss_type_v}') criterion_c = get_criterion(conf.loss_type_c, weights=train_dataset.get_class_weights_c(), rate=conf.ohem_rate) workspace.log(f'Loss type (c): {conf.loss_type_c}') if conf.loss_type_feat_g != 'none': assert isinstance( model, (M.BengaliResNet34V3, M.BengaliResNet34V4, M.BengaliResNet34AGeMV4, M.BengaliSEResNeXt50V4, M.BengaliEfficientNetB0V4, M.BengaliEfficientNetB3V4)) criterion_feat_g = get_criterion(conf.loss_type_feat_g, dim=model.multihead.head_g.dim, n_class=168, s=conf.af_scale_g) workspace.log(f'Loss type (fg): {conf.loss_type_feat_g}') if conf.loss_type_feat_g in ('af', ): sub_models.append(criterion_feat_g) workspace.log('Add criterion_feat_g to sub model') else: criterion_feat_g = None if conf.loss_type_feat_v != 'none': assert isinstance( model, (M.BengaliResNet34V3, M.BengaliResNet34V4, M.BengaliResNet34AGeMV4, M.BengaliSEResNeXt50V4, M.BengaliEfficientNetB0V4, M.BengaliEfficientNetB3V4)) criterion_feat_v = get_criterion(conf.loss_type_feat_v, dim=model.multihead.head_v.dim, n_class=11, s=conf.af_scale_v) workspace.log(f'Loss type (fv): {conf.loss_type_feat_v}') if conf.loss_type_feat_v in ('af', ): sub_models.append(criterion_feat_v) workspace.log('Add criterion_feat_v to sub model') else: criterion_feat_v = None if conf.loss_type_feat_c != 'none': assert isinstance( model, (M.BengaliResNet34V3, M.BengaliResNet34V4, M.BengaliResNet34AGeMV4, M.BengaliSEResNeXt50V4, M.BengaliEfficientNetB0V4, M.BengaliEfficientNetB3V4)) criterion_feat_c = get_criterion(conf.loss_type_feat_c, dim=model.multihead.head_c.dim, n_class=7, s=conf.af_scale_c) workspace.log(f'Loss type (fc): {conf.loss_type_feat_c}') if conf.loss_type_feat_c in ('af', ): sub_models.append(criterion_feat_c) workspace.log('Add criterion_feat_c to sub model') else: criterion_feat_c = None if conf.use_grapheme_code: workspace.log('Use grapheme code classifier') grapheme_classifier = nn.Sequential(nn.BatchNorm1d(168 + 11 + 7), nn.Linear(168 + 11 + 7, 1295)) grapheme_classifier = grapheme_classifier.cuda() grapheme_classifier.train() sub_models.append(grapheme_classifier) criterion_grapheme = L.OHEMCrossEntropyLoss().cuda() else: grapheme_classifier = None criterion_grapheme = None parameters = [{'params': model.parameters()}] + \ [{'params': sub_model.parameters()} for sub_model in sub_models] if conf.optimizer_type == 'adam': optimizer = torch.optim.Adam(parameters, lr=conf.lr) elif conf.optimizer_type == 'sgd': optimizer = torch.optim.SGD(parameters, lr=conf.lr, momentum=0.9, weight_decay=1e-4) elif conf.optimizer_type == 'ranger': optimizer = Ranger(parameters, lr=conf.lr, weight_decay=1e-4) elif conf.optimizer_type == 'radam': optimizer = RAdam(parameters, lr=conf.lr, weight_decay=1e-4) else: raise ValueError(conf.optimizer_type) workspace.log(f'Optimizer type: {conf.optimizer_type}') if conf.use_apex: workspace.log('Apex initialization') _models, optimizer = amp.initialize([model] + sub_models, optimizer, opt_level=conf.apex_opt_level) if len(_models) == 1: model = _models[0] else: model = _models[0] criterion_feat_g = _models[1] criterion_feat_v = _models[2] criterion_feat_c = _models[3] workspace.log('Initialized by Apex') workspace.log(f'{optimizer.__class__.__name__}') for m in _models: workspace.log(f'{m.__class__.__name__}') if conf.scheduler_type == 'cosanl': scheduler = CosineLRWithRestarts( optimizer, conf.batch_size, len(train_dataset), restart_period=conf.cosanl_restart_period, t_mult=conf.cosanl_t_mult) workspace.log(f'restart_period={scheduler.restart_period}') workspace.log(f't_mult={scheduler.t_mult}') elif conf.scheduler_type == 'rop': scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, patience=conf.rop_patience, mode='max', factor=conf.rop_factor, min_lr=1e-6, verbose=True) else: raise ValueError(conf.scheduler_type) train(model, train_loader, val_loader, optimizer, criterion_g, criterion_v, criterion_c, criterion_feat_g, criterion_feat_v, criterion_feat_c, workspace, scheduler=scheduler, n_epoch=conf.n_epoch, cutmix_prob=conf.cutmix_prob, mixup_prob=conf.mixup_prob, freeze_bn_epochs=conf.freeze_bn_epochs, feat_loss_weight=conf.feat_loss_weight, use_apex=conf.use_apex, decrease_ohem_rate=conf.decrease_ohem_rate, use_grapheme_code=conf.use_grapheme_code, grapheme_classifier=grapheme_classifier, criterion_grapheme=criterion_grapheme, final_ft=conf.final_ft)