| import torch | |
| class TemporalMemoryController: | |
| def __init__( | |
| self, | |
| enabled=False, | |
| keyframe_interval=1, | |
| updater='reuse', | |
| updater_cfg=None, | |
| train_with_teacher=False, | |
| distill_weight=1.0, | |
| inference_teacher_forbidden=True, | |
| profiler=None, | |
| debug=False, | |
| ): | |
| if keyframe_interval < 1: | |
| raise ValueError('keyframe_interval must be >= 1') | |
| self.enabled = enabled | |
| self.keyframe_interval = keyframe_interval | |
| self.updater_name = updater.get('type', updater.get('name', 'reuse')) if isinstance(updater, dict) else updater | |
| self.train_with_teacher = train_with_teacher | |
| self.inference_teacher_forbidden = inference_teacher_forbidden | |
| self.debug = debug | |
| self.profiler = profiler | |
| self.frame_index = 0 | |
| self.scene_token = None | |
| self.memory = None | |
| self.last_distillation_loss = None | |
| from .updaters import build_memory_updater | |
| self.updater = build_memory_updater(updater, **(updater_cfg or {})) | |
| from .distillation import DistillationLoss | |
| self.distillation_loss = DistillationLoss(weight=distill_weight) | |
| def reset(self, reason='manual'): | |
| if self.debug: | |
| print(f'[temporal_memory] reset reason={reason}') | |
| self.frame_index = 0 | |
| self.memory = None | |
| def get_or_update( | |
| self, | |
| extract_full, | |
| scene_token=None, | |
| prev_exists=None, | |
| partial_current=None, | |
| ego_motion_embedding=None, | |
| extract_teacher=None, | |
| training=False, | |
| ): | |
| self.last_distillation_loss = None | |
| if not self.enabled: | |
| return extract_full() | |
| if scene_token != self.scene_token: | |
| self.scene_token = scene_token | |
| self.reset(reason='scene_boundary') | |
| if prev_exists is not None and not bool(prev_exists): | |
| self.reset(reason='discontinuity') | |
| is_keyframe = ( | |
| self.memory is None | |
| or self.keyframe_interval == 1 | |
| or self.frame_index % self.keyframe_interval == 0 | |
| ) | |
| if is_keyframe: | |
| if self.debug: | |
| print(f'[temporal_memory] frame={self.frame_index} decision=keyframe') | |
| self.memory = extract_full() | |
| output = self.memory | |
| else: | |
| if self.debug: | |
| print(f'[temporal_memory] frame={self.frame_index} decision=non_keyframe updater={self.updater_name}') | |
| if self.profiler is None: | |
| output = self.updater( | |
| self.memory, | |
| partial_current=partial_current, | |
| ego_motion_embedding=ego_motion_embedding) | |
| else: | |
| with self.profiler.measure('updater'): | |
| output = self.updater( | |
| self.memory, | |
| partial_current=partial_current, | |
| ego_motion_embedding=ego_motion_embedding) | |
| if not training and extract_teacher is not None and self.inference_teacher_forbidden: | |
| raise AssertionError('Full-feature teacher is forbidden on non-keyframe inference.') | |
| if training and self.train_with_teacher and extract_teacher is not None: | |
| with torch.no_grad(): | |
| teacher = extract_teacher() | |
| self.last_distillation_loss = self.distillation_loss(output, teacher) | |
| self.frame_index += 1 | |
| return output | |