def main():
parser = argparse.ArgumentParser()
arg = parser.add_argument
arg('--batch-size', type=int, default=32)
arg('--lr', type=float, default=2e-3)
arg('--workers', type=int, default=4)
arg('--epochs', type=int, default=5)
arg('--mixup-alpha', type=float, default=0)
arg('--arch', type=str, default='seresnext50')
arg('--amp', type=str, default='')
arg('--debug', action='store_true')
args = parser.parse_args()
train_dir = DATA_ROOT / 'train'
valid_dir = DATA_ROOT / 'val'
use_cuda = cuda.is_available()
if args.arch == 'seresnext50':
model = get_seresnet_model(arch="se_resnext50_32x4d",
n_classes=N_CLASSES,
pretrained=False)
elif args.arch == 'seresnext101':
model = get_seresnet_model(arch="se_resnext101_32x4d",
n_classes=N_CLASSES,
pretrained=False)
elif args.arch.startswith("densenet"):
model = get_densenet_model(arch=args.arch)
else:
raise ValueError("No such model")
if use_cuda:
model = model.cuda()
criterion = MixUpSoftmaxLoss(nn.CrossEntropyLoss())
(CACHE_DIR / 'params.json').write_text(
json.dumps(vars(args), indent=4, sort_keys=True))
df_train, class_map = build_dataframe_from_folder(train_dir)
df_valid = build_dataframe_from_folder(valid_dir, class_map)
train_loader = make_loader(args,
TrainDataset,
df_train,
train_transform,
drop_last=True,
shuffle=True)
valid_loader = make_loader(args,
TrainDataset,
df_valid,
test_transform,
shuffle=False)
print(f'{len(train_loader.dataset):,} items in train, '
f'{len(valid_loader.dataset):,} in valid')
train_from_scratch(args, model, train_loader, valid_loader, criterion)
def get_model(arch: str):
if arch == 'seresnext50':
return get_seresnet_model(arch="se_resnext50_32x4d",
n_classes=N_CLASSES,
pretrained=False)
elif arch == 'seresnext101':
return get_seresnet_model(arch="se_resnext101_32x4d",
n_classes=N_CLASSES,
pretrained=False)
elif arch.startswith("densenet"):
return get_densenet_model(arch=arch)
elif arch.startswith("efficientnet"):
return get_efficientnet_model(arch=arch, pretrained=False)
else:
raise ValueError("No such model")