swim/bo-ptycho
1
0
Fork 0
mirror of https://github.com/c-sooyoung/bo-ptycho.git synced 2026-09-17 16:09:09 +09:00
Code Issues Packages Projects Releases Wiki Activity

multi-gpu batched bo

Browse Source
This commit is contained in:
swim
2026-08-12 22:07:45 +09:00
parent 9c66bc84ad
commit 1c85c324c1
6 changed files with 253 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config.yaml
+14 -13
View File
@@ -1,11 +1,11 @@
job:
type: "test"
random_iters: 5
sobo_iters: 10
type: "random+sobo"
random_iters: 8
sobo_iters: 128
io:
input_data_path: '/home/swim/shared/Si_project/data/Si2V1_1.mat'
result_dir: '/home/swim/bo-ptycho/results/test'
input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
result_dir: '/home/swim/bo-ptycho/results/260812-Si/Si2V1_2'
verbosity: 1
ptycho:
@@ -16,7 +16,7 @@ ptycho:
alpha_max: 30
defocus: -200
rot_ang: 0.3
Nlayers: 5
Nlayers: 20
thickness: 250
rbf: 37
Nprobe: 1
@@ -25,8 +25,10 @@ ptycho:
tilt_x: 2
tilt_y: 0
scan_step_size: 0.36
Niter: 5
Niter_save_results: 5
Niter: 100
Niter_save_results: 100
CBED_size: 192
ADU: 1
extra_print_info: ''
@@ -36,11 +38,10 @@ ptycho:
diff_pattern_blur: 1
probe_change_start: 1
object_change_start: 1
grouping: inf
grouping: 512
probe_position_search: 1
regularize_layers: 0.2
variable_probe: 'false'
verbosity
bo:
batch: 4
@@ -50,12 +51,12 @@ bo:
params:
alpha_max:
defocus:
# radius: 100
radius: 100
rot_ang:
Nlayers:
radius: 2
radius: 5
type: int
thickness:
# radius: 100
radius: 100
train_x:
train_y:
job.sub
+1 -1
View File
@@ -3,7 +3,7 @@
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:rtx-6000ada:1
#SBATCH --gres=gpu:rtx-6000ada:4
#SBATCH --time=100:00:00
#SBATCH --output=/home/swim/slurm-logs/job_%j.log
pipelines/__init__.py
+2 -1
View File
@@ -1,5 +1,6 @@
from .sobo import sobo_pipeline
# from .sobo import sobo_pipeline # depreacated
from .test import test_pipeline
from .batched_sobo import sobo_pipeline
job_types = {
'random+sobo': sobo_pipeline,
pipelines/batched_sobo.py
+190
View File
@@ -0,0 +1,190 @@
import os
import traceback
import multiprocessing as mp
import shutil
import bo
import ptycho
def run_ptycho_worker(
worker_id,
gpu_token,
job_config,
metric,
run_id,
result_queue,
):
# This process, and MATLAB launched from it,
# can see exactly one GPU.
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_token
try:
PTYCHOENGINE = ptycho.engines[job_config["ptycho"]["engine"]]
ptycho_engine = PTYCHOENGINE(job_config)
ptycho_engine.run(run_id=run_id)
y_value = ptycho_engine.metric(metric)
result_queue.put(
(worker_id, y_value, None)
)
except Exception:
result_queue.put(
(worker_id, None, traceback.format_exc())
)
def run_batch(
ctx,
gpu_tokens,
job_configs,
metric,
iteration,
):
result_queue = ctx.Queue()
processes = []
for i, job_config in enumerate(job_configs):
# Important: unique run_id for simultaneous jobs.
run_id = f"bo-{iteration:03d}-{i:02d}"
p = ctx.Process(
target=run_ptycho_worker,
args=(
i,
gpu_tokens[i],
job_config,
metric,
run_id,
result_queue,
),
)
p.start()
processes.append(p)
# Four jobs are now running concurrently.
results = [
result_queue.get()
for _ in processes
]
# Synchronization barrier.
for p in processes:
p.join()
# Completion order is arbitrary.
results.sort(key=lambda x: x[0])
for worker_id, _, error in results:
if error is not None:
raise RuntimeError(
f"Ptycho worker {worker_id} failed:\n{error}"
)
return [y_value for _, y_value, _ in results]
def sobo_pipeline(config):
result_dir = config["io"]["result_dir"]
if os.path.exists(result_dir):
shutil.rmtree(result_dir)
os.makedirs(result_dir, exist_ok=True)
RANDOM_ITERS = config["job"].get("random_iters", 0)
SOBO_ITERS = config["job"].get("sobo_iters", 0)
METRIC = config["bo"]["metric"]
BO_BATCH = config["bo"]["batch"]
# SLURM should expose the four GPUs allocated to this job.
visible = os.environ.get("CUDA_VISIBLE_DEVICES")
if visible is None:
raise RuntimeError(
"CUDA_VISIBLE_DEVICES is not set"
)
gpu_tokens = [
token.strip()
for token in visible.split(",")
if token.strip()
]
if len(gpu_tokens) < BO_BATCH:
raise RuntimeError(
f"Expected {BO_BATCH} allocated GPUs, got {len(gpu_tokens)}"
)
# Explicitly use spawn for CUDA / MATLAB isolation.
ctx = mp.get_context("spawn")
# ---------------------------------------------------------
# Random sampling
# ---------------------------------------------------------
randombo = bo.RandomBOEngine(config)
for j in range(RANDOM_ITERS):
print(
f"RANDOM sampling; iteration {j}",
flush=True,
)
job_configs = randombo.ask(n=BO_BATCH)
y_values = run_batch(
ctx=ctx,
gpu_tokens=gpu_tokens,
job_configs=job_configs,
metric=METRIC,
iteration=j,
)
# Only the parent touches BO state / train_x / train_y.
for job_config, y_value in zip(
job_configs,
y_values,
):
randombo.tell(
job_config,
y_value,
)
# ---------------------------------------------------------
# SOBO
# ---------------------------------------------------------
sobo = bo.SingleObjectiveBOEngine(config)
sobo.train_x = randombo.train_x
sobo.train_y = randombo.train_y
for j in range(SOBO_ITERS):
iteration = RANDOM_ITERS + j
print(
f"SOBO sampling; iteration {iteration}",
flush=True,
)
job_configs = sobo.ask(n=BO_BATCH)
y_values = run_batch(
ctx=ctx,
gpu_tokens=gpu_tokens,
job_configs=job_configs,
metric=METRIC,
iteration=iteration,
)
for job_config, y_value in zip(
job_configs,
y_values,
):
sobo.tell(
job_config,
y_value,
)
pipelines/sobo.py
+39 -37
View File
@@ -1,47 +1,49 @@
import os
import bo
import ptycho
# depreacated, use pipelines.batched_sobo.sobo_pipeline()
def sobo_pipeline(config):
# import os
# import bo
# import ptycho
result_dir = config['io']['result_dir']
os.makedirs(result_dir, exist_ok=True)
# def sobo_pipeline(config):
RANDOM_ITERS = config['job'].get('random_iters', 0)
SOBO_ITERS = config['job'].get('sobo_iters')
METRIC = config['bo']['metric']
PTYCHOENGINE = ptycho.engines[config['ptycho']['engine']]
# result_dir = config['io']['result_dir']
# os.makedirs(result_dir, exist_ok=True)
# RANDOM_ITERS = config['job'].get('random_iters', 0)
# SOBO_ITERS = config['job'].get('sobo_iters')
# METRIC = config['bo']['metric']
# PTYCHOENGINE = ptycho.engines[config['ptycho']['engine']]
randombo = bo.RandomBOEngine(config)
# randombo = bo.RandomBOEngine(config)
bo_txt = os.path.join(result_dir, "bo.txt")
with open(bo_txt, "w") as f:
f.write(f" iter\tmetric\t{"\t".join([p[:7] for p in randombo.params])}\n")
# bo_txt = os.path.join(result_dir, "bo.txt")
# with open(bo_txt, "w") as f:
# f.write(f" iter\tmetric\t{"\t".join([p[:7] for p in randombo.params])}\n")
for j in range(RANDOM_ITERS):
print(f"RANDOM sampling; iteration {j}")
job_config = randombo.ask()
ptycho_engine = PTYCHOENGINE(job_config)
ptycho_engine.run(run_id=f"bo-{j:03d}")
y_value = ptycho_engine.metric(METRIC)
randombo.tell(job_config, y_value)
with open(bo_txt, "a") as f:
p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in randombo.params]
f.write(f"{j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")
# for j in range(RANDOM_ITERS):
# print(f"RANDOM sampling; iteration {j}")
# job_config = randombo.ask()
# ptycho_engine = PTYCHOENGINE(job_config)
# ptycho_engine.run(run_id=f"bo-{j:03d}")
# y_value = ptycho_engine.metric(METRIC)
# randombo.tell(job_config, y_value)
# with open(bo_txt, "a") as f:
# p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in randombo.params]
# f.write(f"{j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")
sobo = bo.SingleObjectiveBOEngine(config)
sobo.train_x = randombo.train_x
sobo.train_y = randombo.train_y
# sobo = bo.SingleObjectiveBOEngine(config)
# sobo.train_x = randombo.train_x
# sobo.train_y = randombo.train_y
for j in range(SOBO_ITERS):
print(f"SOBO sampling; iteration {RANDOM_ITERS + j}")
job_config = sobo.ask()
ptycho_engine = PTYCHOENGINE(job_config)
ptycho_engine.run(run_id=f"bo-{RANDOM_ITERS + j:03d}")
y_value = ptycho_engine.metric(METRIC)
sobo.tell(job_config, y_value)
with open(bo_txt, "a") as f:
p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in sobo.params]
f.write(f"{RANDOM_ITERS + j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")
# for j in range(SOBO_ITERS):
# print(f"SOBO sampling; iteration {RANDOM_ITERS + j}")
# job_config = sobo.ask()
# ptycho_engine = PTYCHOENGINE(job_config)
# ptycho_engine.run(run_id=f"bo-{RANDOM_ITERS + j:03d}")
# y_value = ptycho_engine.metric(METRIC)
# sobo.tell(job_config, y_value)
# with open(bo_txt, "a") as f:
# p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in sobo.params]
# f.write(f"{RANDOM_ITERS + j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")
ptycho/fold_slice.py
+7 -4
View File
@@ -11,9 +11,7 @@ class FoldSlicePtychoEngine(PtychoEngine):
def __init__(self, config):
super().__init__(config)
self._output_dir = os.path.join(config['io']['result_dir'], 'fold_slice')
self._fold_slice_path = self.config['ptycho']['path']
self._setup_txt_path = os.path.join(self._output_dir, 'setup.txt')
self.metric_methods = {
'log_fourier': self._log_fourier_metric,
}
@@ -21,6 +19,9 @@ class FoldSlicePtychoEngine(PtychoEngine):
def run(self, run_id="") -> None:
self._output_dir = os.path.join(self.config['io']['result_dir'], f'fold_slice-{run_id}')
self._setup_txt_path = os.path.join(self._output_dir, 'setup.txt')
# generate setup.txt for fold_slice input
fold_slice_dict = {}
fold_slice_dict['raw_data'] = self.config['io']['input_data_path']
@@ -82,9 +83,11 @@ class FoldSlicePtychoEngine(PtychoEngine):
log_fourier_error = self._log_fourier_metric()
# os.makedirs(os.path.join(self.config['io']['result_dir'], "mat"), exist_ok=True)
os.makedirs(os.path.join(self.config['io']['result_dir'], "tiff"), exist_ok=True)
# os.makedirs(os.path.join(self.config['io']['result_dir'], "tiff"), exist_ok=True)
# shutil.copy(mat_path, os.path.join(self.config['io']['result_dir'], "mat", f"{log_fourier_error:.4f}_{run_id}.mat")) # saving .mat files takes a lot of space (expect 20+ GB for 64*64 scan size, 300 iterations)
shutil.copy(image_path, os.path.join(self.config['io']['result_dir'], "tiff", f"{log_fourier_error:.4f}_{run_id}.tiff"))
# shutil.copy(image_path, os.path.join(self.config['io']['result_dir'], "tiff", f"{log_fourier_error:.4f}_{run_id}.tiff"))
shutil.rmtree(self._output_dir)
def metric(self, names):