Shape-as-Point/optim_hierarchy.py

81 lines
2.7 KiB
Python
Raw Normal View History

2021-11-08 10:09:50 +00:00
import argparse
2023-05-26 12:59:53 +00:00
import os
2021-11-08 10:09:50 +00:00
from src.utils import load_config
2023-05-26 12:59:53 +00:00
os.environ["MKL_THREADING_LAYER"] = "GNU"
2021-11-08 10:09:50 +00:00
2023-05-26 12:59:53 +00:00
def main():
parser = argparse.ArgumentParser(description="MNIST toy experiment")
parser.add_argument("config", type=str, help="Path to config file.")
parser.add_argument("--start_res", type=int, default=-1, help="Resolution to start with.")
parser.add_argument("--object_id", type=int, default=-1, help="Object index.")
2021-11-08 10:09:50 +00:00
2023-05-26 12:59:53 +00:00
args, unknown = parser.parse_known_args()
cfg = load_config(args.config, "configs/default.yaml")
2021-11-08 10:09:50 +00:00
2023-05-26 12:59:53 +00:00
resolutions = [32, 64, 128, 256]
iterations = [1000, 1000, 1000, 200]
lrs = [2e-3, 2e-3 * 0.7, 2e-3 * (0.7**2), 2e-3 * (0.7**3)] # reduce lr
for idx, (res, iteration, lr) in enumerate(zip(resolutions, iterations, lrs)):
if res < args.start_res:
2021-11-08 10:09:50 +00:00
continue
2023-05-26 12:59:53 +00:00
if res > cfg["model"]["grid_res"]:
2021-11-08 10:09:50 +00:00
continue
2023-05-26 12:59:53 +00:00
psr_sigma = 2 if res <= 128 else 3
2021-11-08 10:09:50 +00:00
if res > 128:
2023-05-26 12:59:53 +00:00
psr_sigma = 5 if "thingi_noisy" in args.config else 3
2021-11-08 10:09:50 +00:00
if args.object_id != -1:
2023-05-26 12:59:53 +00:00
out_dir = os.path.join(cfg["train"]["out_dir"], "object_%02d" % args.object_id, "res_%d" % res)
2021-11-08 10:09:50 +00:00
else:
2023-05-26 12:59:53 +00:00
out_dir = os.path.join(cfg["train"]["out_dir"], "res_%d" % res)
2021-11-08 10:09:50 +00:00
# sample from mesh when resampling is enabled, otherwise reuse the pointcloud
2023-05-26 12:59:53 +00:00
init_shape = "mesh" if cfg["train"]["resample_every"] > 0 else "pointcloud"
2021-11-08 10:09:50 +00:00
if args.object_id != -1:
2023-05-26 12:59:53 +00:00
input_mesh = (
"None"
if idx == 0
else os.path.join(
cfg["train"]["out_dir"],
"object_%02d" % args.object_id,
"res_%d" % (resolutions[idx - 1]),
"vis",
init_shape,
"%04d.ply" % (iterations[idx - 1]),
)
)
2021-11-08 10:09:50 +00:00
else:
2023-05-26 12:59:53 +00:00
input_mesh = (
"None"
if idx == 0
else os.path.join(
cfg["train"]["out_dir"],
"res_%d" % (resolutions[idx - 1]),
"vis",
init_shape,
"%04d.ply" % (iterations[idx - 1]),
)
)
cmd = "export MKL_SERVICE_FORCE_INTEL=1 && "
cmd += (
"python optim.py %s --model:grid_res %d --model:psr_sigma %d \
2021-11-08 10:09:50 +00:00
--train:input_mesh %s --train:total_epochs %d \
--train:out_dir %s --train:lr_pcl %f \
2023-05-26 12:59:53 +00:00
--data:object_id %d"
% (args.config, res, psr_sigma, input_mesh, iteration, out_dir, lr, args.object_id)
)
2021-11-08 10:09:50 +00:00
print(cmd)
os.system(cmd)
2023-05-26 12:59:53 +00:00
if __name__ == "__main__":
2021-11-08 10:09:50 +00:00
main()