def _init_cuda(self):
if self.use_cuda_:
torch.cuda.init()
from falkon.cuda import initialization
initialization.init(self._base_opt)
self.num_gpus = devices.num_gpus(self.options)
t_s = time.time()
fn(Ac)
t_e = time.time()
timings.append(t_e - t_s)
print("Exp %s - N %d - Rep %d - %.2fs" %
(exp, N, j, timings[-1]),
flush=True)
if exp['torch']:
torch.cuda.empty_cache()
exp['timings'].append(min(timings))
return experiments
if __name__ == "__main__":
init_opt = falkon.FalkonOptions(compute_arch_speed=False)
initialization.init(init_opt)
gpu_info = [
v for k, v in devices.get_device_info(init_opt).items() if k >= 0
]
num_gpu = len(gpu_info)
defaultN32 = [
10_000, 20_000, 30_000, 40_000, 50_000, 65_000, 80_000, 100_000,
120_000, 140_000
]
#defaultN64 = [10_000, 20_000, 30_000, 40_000, 50_000, 65_000, 80_000]
falkon.FalkonOptions(chol_force_ooc=True,
chol_par_blk_multiplier=2,
compute_arch_speed=False)
experiments = [
{
elif exp_num == 2:
opt = dataclasses.replace(opt, cpu_preconditioner=True, keops_active="no")
dtype = DataType.float32
elif exp_num == 3:
opt = dataclasses.replace(opt, cpu_preconditioner=False, keops_active="no")
dtype = DataType.float32
elif exp_num == 4:
opt = dataclasses.replace(opt, cpu_preconditioner=False, keops_active="no")
dtype = DataType.float32
elif exp_num == 5:
opt = dataclasses.replace(opt, cpu_preconditioner=False, keops_active="force")
dtype = DataType.float32
else:
raise ValueError("exp num %d not valid" % (exp_num))
data = load_data(dset, data_type=dtype)
initialization.init(opt)
print("\n\n --- Running Experiment %d -- %s" % (exp_num, opt))
data = list(data)
data[0] = data[0].pin_memory()
data[1] = data[1].pin_memory()
data[2] = data[2].pin_memory()
data[3] = data[3].pin_memory()
t_s = time.time()
flk = run_single(dset, data[0], data[1], data[2], data[3], data[4], show_intermediate_errors, opt, params)
t_e = time.time()
print("Timing for Experiment %d: %s -- fit times %s" % (exp_num, t_e - t_s, flk.fit_times_))
def load_data(dset, data_type):
load_fn = get_load_fn(dset)
return load_fn(dtype=data_type.to_numpy_dtype(), as_torch=True)
def initialize_cuda():
# your setup code goes here, executed ahead of first test
opt = BaseOptions(compute_arch_speed=False, use_cpu=False)
if decide_cuda():
initialization.init(opt)