|
|
import os |
|
|
import time |
|
|
from dataclasses import dataclass |
|
|
|
|
|
from configs.mode import FaceSwapMode |
|
|
from configs.singleton import Singleton |
|
|
|
|
|
|
|
|
@Singleton |
|
|
@dataclass |
|
|
class TrainConfig: |
|
|
mode = FaceSwapMode.MANY_TO_MANY |
|
|
source_name: str = "" |
|
|
|
|
|
dataset_index: str = "/data/dataset/faceswap/full.pkl" |
|
|
dataset_root: str = "/data/dataset/faceswap" |
|
|
|
|
|
batch_size: int = 8 |
|
|
num_threads: int = 8 |
|
|
same_rate: float = 0.5 |
|
|
lr: float = 5e-5 |
|
|
grad_clip: float = 1000.0 |
|
|
|
|
|
use_ddp: bool = True |
|
|
|
|
|
mouth_mask: bool = True |
|
|
eye_hm_loss: bool = False |
|
|
mouth_hm_loss: bool = False |
|
|
|
|
|
load_checkpoint = None |
|
|
|
|
|
identity_extractor_config = { |
|
|
"f_3d_checkpoint_path": "/checkpoints/Deep3DFaceRecon/epoch_20_new.pth", |
|
|
"f_id_checkpoint_path": "/checkpoints/arcface/ms1mv3_arcface_r100_fp16_backbone.pth", |
|
|
"bfm_folder": "/checkpoints/useful_ckpt/BFM", |
|
|
"hrnet_path": "/checkpoints/useful_ckpt/face_98lmks/HR18-WFLW.pth", |
|
|
} |
|
|
|
|
|
visualize_interval: int = 100 |
|
|
plot_interval: int = 100 |
|
|
max_iters: int = 1000000 |
|
|
checkpoint_interval: int = 40000 |
|
|
|
|
|
exp_name: str = "exp_base" |
|
|
log_basedir: str = "/data/logs/hififace/" |
|
|
checkpoint_basedir = "/data/checkpoints/hififace" |
|
|
|
|
|
def __post_init__(self): |
|
|
time_stamp = int(time.time() * 1000) |
|
|
self.log_dir = os.path.join(self.log_basedir, f"{self.exp_name}_{time_stamp}") |
|
|
self.checkpoint_dir = os.path.join(self.checkpoint_basedir, f"{self.exp_name}_{time_stamp}") |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
tc = TrainConfig() |
|
|
print(tc.log_dir) |
|
|
|
