Basic Framework

This commit is contained in:
hofee
2024-08-18 00:37:17 +08:00
commit 73dcd592df
14 changed files with 733 additions and 0 deletions

74
configs/config.py Normal file
View File

@@ -0,0 +1,74 @@
import argparse
import os.path
import shutil
import yaml
class ConfigManager:
config = None
config_path = None
@staticmethod
def get(*args):
result = ConfigManager.config
for arg in args:
result = result[arg]
return result
@staticmethod
def load_config_with(config_file_path):
ConfigManager.config_path = config_file_path
if not os.path.exists(ConfigManager.config_path):
raise ValueError(f"Config file <{config_file_path}> does not exist")
with open(config_file_path, 'r') as file:
ConfigManager.config = yaml.safe_load(file)
@staticmethod
def backup_config_to(target_config_dir, file_name, prefix="config"):
file_name = f"{prefix}_{file_name}.yaml"
target_config_file_path = str(os.path.join(target_config_dir, file_name))
shutil.copy(ConfigManager.config_path, target_config_file_path)
@staticmethod
def load_config():
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default='', help='config file path')
args = parser.parse_args()
if args.config:
ConfigManager.load_config_with(args.config)
@staticmethod
def print_config(key: str = None, group: dict = None, level=0):
table_size = 80
if key and group:
value = group[key]
if type(value) is dict:
print("\t" * level + f"+-{key}:")
for k in value:
ConfigManager.print_config(k, value, level=level + 1)
else:
print("\t" * level + f"| {key}: {value}")
elif key:
ConfigManager.print_config(key, ConfigManager.config, level=level)
else:
print("+" + "-" * table_size + "+")
print(f"| Configurations in <{ConfigManager.config_path}>:")
print("+" + "-" * table_size + "+")
for key in ConfigManager.config:
ConfigManager.print_config(key, level=level + 1)
print("+" + "-" * table_size + "+")
''' ------------ Debug ------------ '''
if __name__ == "__main__":
test_args = ['--config', r'configs\train_config.yaml']
test_parser = argparse.ArgumentParser()
test_parser.add_argument('--config', type=str, default='', help='config file path')
test_args = test_parser.parse_args(test_args)
if test_args.config:
ConfigManager.load_config_with(test_args.config)
ConfigManager.print_config()
print()
pipeline = ConfigManager.get('settings', 'train', "dataset", 'batch_size')
ConfigManager.print_config('settings')
print(pipeline)

73
configs/train_config.yaml Normal file
View File

@@ -0,0 +1,73 @@
# Train config file
settings:
general:
seed: 0
device: cuda
cuda_visible_devices: "0,1,2,3,4,5,6,7"
parallel: True
experiment:
name: train_experiment
root_dir: "experiments"
use_checkpoint: True
epoch: -1 # -1 stands for last epoch
max_epochs: 5000
save_checkpoint_interval: 1
test_first: True
train:
optimizer:
type: adam
lr: 0.0001
losses: # loss type : weight
gf_loss: 1.0
dataset:
name: synthetic_train_train_dataset
source: nbv1
data_type: train
ratio: 0.1
batch_size: 128
num_workers: 96
test:
batch_size: 16
frequency: 3
dataset_list:
- name: synthetic_test_train_dataset
source: nbv1
data_type: train
eval_list:
ratio: 0.00001
batch_size: 16
num_workers: 16
pipeline:
pts_encoder: pointnet
view_finder: gradient_field
datasets:
general:
data_dir: "../data"
modules:
general:
pts_channels: 3
feature_dim: 1024
per_point_feature: False
pts_encoder:
pointnet:
pointnet++:
params_name: light
pointnet++rgb:
params_name: light
target_layer: 3
rgb_feat_dim: 384
view_finder:
gradient_field:
pose_mode: rot_matrix
regression_head: Rx_Ry
sample_mode: ode
sample_repeat: 50
sampling_steps: 500
sde_mode: ve