File size: 3,530 Bytes
6a176cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | 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
|