os.makedirs(save_dir)
# Create the data loaders
train_loader, valid_loader, test_loader = imagenet_subset.get_dataloaders(
batch_size=args.batch_size, data_path=args.data_path, num_workers=4)
# Load and initialize the model
params = parse_config.load_config('../../cfg/alexnet_ckn.cfg')
num_layers = len(params['num_filters'])
if args.num_filters > 0:
params['num_filters'] = [args.num_filters] * num_layers
params['patch_sigma'] = [bw] * num_layers
params['patch_kernel'] = [args.kernel] * num_layers
params['matern_order'] = [args.matern_order] * num_layers
layers = parse_config.create_layers(params)
model = net.CKN(layers)
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model).to(defaults.device)
model.module.init(train_loader)
else:
model.to(defaults.device).init(train_loader)
print('Done with the initialization')
# Create the step size schedule
if args.step_size_method == 'fixed':
step_size_schedule = None
update_step_size_method = 'fixed'
elif args.step_size_method == 'best-ckn-128':
def step_size_schedule(x):
bw = np.median(
sklearn.metrics.pairwise.pairwise_distances(
train_loader.dataset.features[0].numpy()).reshape(-1))
print('Bandwidth from median heuristic:')
else:
bw = args.bw
# Load and initialize the model
model_params = parse_config.load_config('../../cfg/kn.cfg')
if args.num_filters > 0:
for key in model_params:
model_params[key] *= args.num_layers
model_params['num_filters'] = [args.num_filters] * args.num_layers
model_params['patch_sigma'] = [bw] * args.num_layers
layers = parse_config.create_layers(model_params)
model = net.CKN(layers).to(defaults.device)
model.init(train_loader)
print('Done with initialization')
# Set up the data, parameters, model, results, and optimizer objects
data = opt_structures.Data(train_labeled_loader, train_unlabeled_loader,
valid_loader, test_loader)
params = opt_structures.Params(batch_size=args.batch_size,
ckn=True,
epsilon=args.epsilon,
eval_test_every=args.eval_test_every,
lambda_cov=args.lambda_cov,
lambda_params=0,
lr=args.lr,
lr_schedule=None,