请问怎么引入新的reference model

[2016110071]

[Jayce-Ping]

目前没有为专门实现update_reference_model的逻辑。不过可以直接用现有接口来手动实现：

Flow-Factory/src/flow_factory/models/abc.py

Lines 596 to 713 in 33fd342

	def add_named_parameters(
	self,
	name: str,
	target_components: Optional[Union[str, List[str]]] = None,
	device: Optional[Union[torch.device, str]] = None,
	overwrite: bool = True,
	) -> None:
	"""
	Store current trainable parameters snapshot under a name.
	Args:
	name: Identifier for this parameter snapshot
	target_modules: Component names to store. Defaults to components with trainable params.
	device: Storage device (defaults to 'cpu')
	overwrite: Whether to overwrite existing snapshot
	"""
	if name in self._named_parameters and not overwrite:
	raise KeyError(f"Named parameters '{name}' exists. Use overwrite=True.")
	# Normalize target_modules - filter only those with trainable params
	if target_components is None:
	target_components = [k for k, v in self.target_module_map.items() if v]
	elif isinstance(target_components, str):
	target_components = [target_components]
	# Validate
	invalid = set(target_components) - set(self.target_module_map.keys())
	if invalid:
	raise ValueError(f"Invalid target_modules: {invalid}. Valid: {list(self.target_module_map.keys())}")
	device = torch.device(device) if device else torch.device('cpu')
	params = self._get_component_parameters(target_components)
	if not params:
	raise ValueError(f"No trainable parameters found in {target_components}")
	self._named_parameters[name] = NamedParametersInfo(
	target_components=target_components,
	ema_wrapper=EMAModuleWrapper(
	parameters=params,
	decay=0.0,
	update_step_interval=0,
	device=device,
	),
	)
	logger.info(f"Stored named parameters '{name}' for {target_components} on {device}")
	@contextmanager
	def use_named_parameters(self, name: str):
	"""
	Context manager to temporarily use named parameters.
	Args:
	name: Name of stored parameters snapshot
	Usage:
	adapter.add_named_parameters('init')
	# ... training ...
	with adapter.use_named_parameters('init'):
	evaluate(model) # Uses stored weights
	# Current weights restored
	"""
	if name not in self._named_parameters:
	raise KeyError(f"'{name}' not found. Available: {self.list_named_parameters()}")
	info = self._named_parameters[name]
	params = self._get_component_parameters(info.target_components)
	with info.ema_wrapper.use_ema_parameters(params):
	yield
	def update_named_parameters(
	self,
	name: str,
	target_components: Optional[Union[str, List[str]]] = None,
	new_parameters: Optional[Iterable[torch.nn.Parameter]] = None,
	) -> None:
	"""
	Update existing named parameters with specified or current values.
	Args:
	name: Name of snapshot to update
	target_modules: Components to update. Defaults to originally stored components.
	new_parameters: Parameters to copy from. Defaults to current model parameters.
	"""
	if name not in self._named_parameters:
	raise KeyError(f"'{name}' not found.")
	info = self._named_parameters[name]
	# Resolve target_modules
	if target_components is None:
	target_components = info.target_components
	elif isinstance(target_components, str):
	target_components = [target_components]
	if not set(target_components).issubset(set(info.target_components)):
	raise ValueError(f"Must be subset of original: {info.target_components}")
	# Resolve parameters
	if new_parameters is None:
	new_parameters = self._get_component_parameters(target_components)
	else:
	new_parameters = list(new_parameters)
	# Validate param count
	if len(new_parameters) != len(info.ema_wrapper.ema_parameters):
	raise ValueError(
	f"Parameter count mismatch: got {len(new_parameters)}, "
	f"expected {len(info.ema_wrapper.ema_parameters)}"
	)
	# Update
	with torch.no_grad():
	for ema_param, param in zip(info.ema_wrapper.ema_parameters, new_parameters, strict=True):
	ema_param.data.copy_(param.detach().to(ema_param.device))
	logger.info(f"Updated named parameters '{name}'")

在自定义的Trainer中用add_named_parameters来手动加入一个参数为moving_ref，然后在需要的时候用update_moving_ref去更新就可以了。