import argparse
from pathlib import Path
import torch
def _extract_state(ckpt, use_ema=False):
if isinstance(ckpt, dict):
if use_ema and 'ema' in ckpt:
model = ckpt['ema']
return model.float().state_dict(), ckpt, 'ema'
if 'model' in ckpt:
model = ckpt['model']
return model.float().state_dict(), ckpt, 'model'
if 'state_dict' in ckpt:
return ckpt['state_dict'], ckpt, 'state_dict'
if isinstance(ckpt, dict):
return ckpt, ckpt, 'state_dict'
raise ValueError('Unsupported checkpoint format')
def average_checkpoints(model_dir, output, pattern='*.pt', use_ema=False, strict_keys=True):
model_dir = Path(model_dir)
files = sorted(model_dir.glob(pattern))
if len(files) < 2:
raise ValueError(f'Need at least 2 checkpoints, got {len(files)} in {model_dir}')
first_state, first_ckpt, first_key = _extract_state(torch.load(files[0], map_location='cpu'), use_ema)
keys = list(first_state.keys())
avg_state = {}
for k in keys:
v = first_state[k]
avg_state[k] = v.float().clone() if torch.is_floating_point(v) else v
for f in files[1:]:
state, _, _ = _extract_state(torch.load(f, map_location='cpu'), use_ema)
if strict_keys and set(state.keys()) != set(keys):
raise ValueError(f'Key mismatch in {f}')
for k in keys:
if k not in state:
if strict_keys:
raise ValueError(f'Missing key {k} in {f}')
continue
v = state[k]
if torch.is_floating_point(v):
avg_state[k] += v.float()
else:
avg_state[k] = avg_state[k]
for k in keys:
if torch.is_floating_point(avg_state[k]):
avg_state[k] /= float(len(files))
if first_key in ('model', 'ema'):
first_ckpt[first_key].load_state_dict(avg_state, strict=False)
if first_key == 'ema' and 'model' in first_ckpt:
first_ckpt['model'].load_state_dict(avg_state, strict=False)
torch.save(first_ckpt, output)
else:
torch.save(avg_state, output)
return files
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--dir', type=str, required=True, help='directory with checkpoints')
parser.add_argument('--output', type=str, required=True, help='output checkpoint path')
parser.add_argument('--pattern', type=str, default='*.pt', help='checkpoint filename pattern')
parser.add_argument('--use-ema', action='store_true', help='average EMA weights if available')
parser.add_argument('--strict-keys', action='store_true', help='require identical state_dict keys')
args = parser.parse_args()
files = average_checkpoints(
model_dir=args.dir,
output=args.output,
pattern=args.pattern,
use_ema=args.use_ema,
strict_keys=args.strict_keys,
)
print(f'Averaged {len(files)} checkpoints to {args.output}')
if name == 'main':
main()
已新增模型参数平均脚本 average_weights.py，输入一个目录，自动把目录下多个 checkpoint 做参数平均并保存为新权重。支持使用 EMA 权重，且可选择严格检查 key 是否一致。
用法示例：
python average_weights.py --dir "E:\同致电子\动态分辨率\runs\train\exp\weights" --output "E:\同致电子\动态分辨率\avg.pt"
如果你要用 EMA 权重平均：
python average_weights.py --dir "E:\...\weights" --output "E:\...\avg.pt" --use-ema
如果想强制所有模型 key 必须完全一致：
python average_weights.py --dir "E:\...\weights" --output "E:\...\avg.pt" --strict-keys
说明：
- 默认匹配模式是 *.pt，可用 --pattern 调整。
- 生成的 avg.pt 会保留与首个 checkpoint 相同的结构（model 或 ema），后续可以直接拿去量化。
如果你希望只平均指定的几个 epoch（不是目录中全部），告诉我你想按什么规则筛选（文件名、epoch 号列表、正则），我可以加筛选参数。
为了使的平均后的模型效果最好，--use-ema和--strict-keys要开启吗
简短结论：
- --use-ema：建议开启，前提是这些 checkpoint 里确实有 EMA。EMA 往往更稳、更接近最优验证效果，平均 EMA 通常比平均普通权重更好。
- --strict-keys：建议开启，前提是这些模型结构完全一致（同一 cfg/类数/分割头）。这样能保证平均的是同一层，避免"误对齐"带来的退化。
例外情况：
- 如果部分模型来自不同结构或类数，--strict-keys 会直接报错，这时要么筛掉不一致的模型，要么不用 strict（但平均效果通常不如严格一致）。
- 如果某些 checkpoint 没有 EMA，--use-ema 会回退到 model 吗？当前脚本不会自动回退，会按 ema 读取不到就报错，因此要确认这些权重都带 EMA。
最稳妥的组合（同结构同数据）：--use-ema --strict-keys。