xsl
ce445b64a6
部署修复： - torch.load 增加 weights_only=False patch，兼容 PyTorch 2.6+ 加载旧权重 - OSS 改为懒加载，本地用 output_format=base64 无需配凭证即可启动 - 补全被 gitignore 误排除的必需代码：core/models/layers/data、models/layers/data、keypoints/lib - webui 训练命令 --xformers 改 --sdpa（修复 xformers 无 CUDA 支持报错） 功能调整： - hair_grow_service 端口改 8899、preview 路由修复（send_file） - list_hairstyles 增加发型白名单，测试页只展示当前5个发型 新增脚本： - train_lora_parallel.py：直接调 kohya 并行训练 LoRA（绕过 photo_service 串行限制） - train_hairstyles_parallel.py / train_batch_stepC.py：批量训练辅助脚本 - scripts/sync_data_to_server.sh：大文件断点续传到云服务器 文档： - docs/换发型集成文档.md：换发型完整流程、服务架构、资源依赖、训练方法、集成步骤
161 lines
3.4 KiB
Python
161 lines
3.4 KiB
Python
|
|
# ------------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft
|
|
# Licensed under the MIT License.
|
|
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
|
# ------------------------------------------------------------------------------
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import os
|
|
|
|
from yacs.config import CfgNode as CN
|
|
|
|
|
|
_C = CN()
|
|
|
|
_C.OUTPUT_DIR = ''
|
|
_C.LOG_DIR = ''
|
|
_C.DATA_DIR = ''
|
|
_C.GPUS = (0,)
|
|
_C.WORKERS = 4
|
|
_C.PRINT_FREQ = 20
|
|
_C.AUTO_RESUME = False
|
|
_C.PIN_MEMORY = True
|
|
_C.RANK = 0
|
|
|
|
# Cudnn related params
|
|
_C.CUDNN = CN()
|
|
_C.CUDNN.BENCHMARK = True
|
|
_C.CUDNN.DETERMINISTIC = False
|
|
_C.CUDNN.ENABLED = True
|
|
|
|
# common params for NETWORK
|
|
_C.MODEL = CN()
|
|
_C.MODEL.NAME = 'pose_hrnet'
|
|
_C.MODEL.INIT_WEIGHTS = True
|
|
_C.MODEL.PRETRAINED = ''
|
|
_C.MODEL.NUM_JOINTS = 17
|
|
_C.MODEL.TAG_PER_JOINT = True
|
|
_C.MODEL.TARGET_TYPE = 'gaussian'
|
|
_C.MODEL.IMAGE_SIZE = [256, 256] # width * height, ex: 192 * 256
|
|
_C.MODEL.HEATMAP_SIZE = [64, 64] # width * height, ex: 24 * 32
|
|
_C.MODEL.SIGMA = 2
|
|
_C.MODEL.EXTRA = CN(new_allowed=True)
|
|
|
|
_C.LOSS = CN()
|
|
_C.LOSS.USE_OHKM = False
|
|
_C.LOSS.TOPK = 8
|
|
_C.LOSS.USE_TARGET_WEIGHT = True
|
|
_C.LOSS.USE_DIFFERENT_JOINTS_WEIGHT = False
|
|
|
|
# DATASET related params
|
|
_C.DATASET = CN()
|
|
_C.DATASET.ROOT = ''
|
|
_C.DATASET.DATASET = 'mpii'
|
|
_C.DATASET.TRAIN_SET = 'train'
|
|
_C.DATASET.TEST_SET = 'valid'
|
|
_C.DATASET.DATA_FORMAT = 'jpg'
|
|
_C.DATASET.HYBRID_JOINTS_TYPE = ''
|
|
_C.DATASET.SELECT_DATA = False
|
|
|
|
# training data augmentation
|
|
_C.DATASET.FLIP = True
|
|
_C.DATASET.SCALE_FACTOR = 0.25
|
|
_C.DATASET.ROT_FACTOR = 30
|
|
_C.DATASET.PROB_HALF_BODY = 0.0
|
|
_C.DATASET.NUM_JOINTS_HALF_BODY = 8
|
|
_C.DATASET.COLOR_RGB = False
|
|
|
|
# train
|
|
_C.TRAIN = CN()
|
|
|
|
_C.TRAIN.LR_FACTOR = 0.1
|
|
_C.TRAIN.LR_STEP = [90, 110]
|
|
_C.TRAIN.LR = 0.001
|
|
|
|
_C.TRAIN.OPTIMIZER = 'adam'
|
|
_C.TRAIN.MOMENTUM = 0.9
|
|
_C.TRAIN.WD = 0.0001
|
|
_C.TRAIN.NESTEROV = False
|
|
_C.TRAIN.GAMMA1 = 0.99
|
|
_C.TRAIN.GAMMA2 = 0.0
|
|
|
|
_C.TRAIN.BEGIN_EPOCH = 0
|
|
_C.TRAIN.END_EPOCH = 140
|
|
|
|
_C.TRAIN.RESUME = False
|
|
_C.TRAIN.CHECKPOINT = ''
|
|
|
|
_C.TRAIN.BATCH_SIZE_PER_GPU = 32
|
|
_C.TRAIN.SHUFFLE = True
|
|
|
|
# testing
|
|
_C.TEST = CN()
|
|
|
|
# size of images for each device
|
|
_C.TEST.BATCH_SIZE_PER_GPU = 32
|
|
# Test Model Epoch
|
|
_C.TEST.FLIP_TEST = False
|
|
_C.TEST.POST_PROCESS = False
|
|
_C.TEST.SHIFT_HEATMAP = False
|
|
|
|
_C.TEST.USE_GT_BBOX = False
|
|
|
|
# nms
|
|
_C.TEST.IMAGE_THRE = 0.1
|
|
_C.TEST.NMS_THRE = 0.6
|
|
_C.TEST.SOFT_NMS = False
|
|
_C.TEST.OKS_THRE = 0.5
|
|
_C.TEST.IN_VIS_THRE = 0.0
|
|
_C.TEST.COCO_BBOX_FILE = ''
|
|
_C.TEST.BBOX_THRE = 1.0
|
|
_C.TEST.MODEL_FILE = ''
|
|
|
|
# debug
|
|
_C.DEBUG = CN()
|
|
_C.DEBUG.DEBUG = False
|
|
_C.DEBUG.SAVE_BATCH_IMAGES_GT = False
|
|
_C.DEBUG.SAVE_BATCH_IMAGES_PRED = False
|
|
_C.DEBUG.SAVE_HEATMAPS_GT = False
|
|
_C.DEBUG.SAVE_HEATMAPS_PRED = False
|
|
|
|
|
|
def update_config(cfg, args):
|
|
cfg.defrost()
|
|
cfg.merge_from_file(args.cfg)
|
|
cfg.merge_from_list(args.opts)
|
|
|
|
if args.modelDir:
|
|
cfg.OUTPUT_DIR = args.modelDir
|
|
|
|
if args.logDir:
|
|
cfg.LOG_DIR = args.logDir
|
|
|
|
if args.dataDir:
|
|
cfg.DATA_DIR = args.dataDir
|
|
|
|
cfg.DATASET.ROOT = os.path.join(
|
|
cfg.DATA_DIR, cfg.DATASET.ROOT
|
|
)
|
|
|
|
cfg.MODEL.PRETRAINED = os.path.join(
|
|
cfg.DATA_DIR, cfg.MODEL.PRETRAINED
|
|
)
|
|
|
|
if cfg.TEST.MODEL_FILE:
|
|
cfg.TEST.MODEL_FILE = os.path.join(
|
|
cfg.DATA_DIR, cfg.TEST.MODEL_FILE
|
|
)
|
|
|
|
cfg.freeze()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
import sys
|
|
with open(sys.argv[1], 'w') as f:
|
|
print(_C, file=f)
|
|
|