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)