Track-Anything / tracker /inference /memory_manager.py
watchtowerss's picture
track-anything --version 1
4d1ebf3
raw
history blame
12.4 kB
import torch
import warnings
from inference.kv_memory_store import KeyValueMemoryStore
from model.memory_util import *
class MemoryManager:
"""
Manages all three memory stores and the transition between working/long-term memory
"""
def __init__(self, config):
self.hidden_dim = config['hidden_dim']
self.top_k = config['top_k']
self.enable_long_term = config['enable_long_term']
self.enable_long_term_usage = config['enable_long_term_count_usage']
if self.enable_long_term:
self.max_mt_frames = config['max_mid_term_frames']
self.min_mt_frames = config['min_mid_term_frames']
self.num_prototypes = config['num_prototypes']
self.max_long_elements = config['max_long_term_elements']
# dimensions will be inferred from input later
self.CK = self.CV = None
self.H = self.W = None
# The hidden state will be stored in a single tensor for all objects
# B x num_objects x CH x H x W
self.hidden = None
self.work_mem = KeyValueMemoryStore(count_usage=self.enable_long_term)
if self.enable_long_term:
self.long_mem = KeyValueMemoryStore(count_usage=self.enable_long_term_usage)
self.reset_config = True
def update_config(self, config):
self.reset_config = True
self.hidden_dim = config['hidden_dim']
self.top_k = config['top_k']
assert self.enable_long_term == config['enable_long_term'], 'cannot update this'
assert self.enable_long_term_usage == config['enable_long_term_count_usage'], 'cannot update this'
self.enable_long_term_usage = config['enable_long_term_count_usage']
if self.enable_long_term:
self.max_mt_frames = config['max_mid_term_frames']
self.min_mt_frames = config['min_mid_term_frames']
self.num_prototypes = config['num_prototypes']
self.max_long_elements = config['max_long_term_elements']
def _readout(self, affinity, v):
# this function is for a single object group
return v @ affinity
def match_memory(self, query_key, selection):
# query_key: B x C^k x H x W
# selection: B x C^k x H x W
num_groups = self.work_mem.num_groups
h, w = query_key.shape[-2:]
query_key = query_key.flatten(start_dim=2)
selection = selection.flatten(start_dim=2) if selection is not None else None
"""
Memory readout using keys
"""
if self.enable_long_term and self.long_mem.engaged():
# Use long-term memory
long_mem_size = self.long_mem.size
memory_key = torch.cat([self.long_mem.key, self.work_mem.key], -1)
shrinkage = torch.cat([self.long_mem.shrinkage, self.work_mem.shrinkage], -1)
similarity = get_similarity(memory_key, shrinkage, query_key, selection)
work_mem_similarity = similarity[:, long_mem_size:]
long_mem_similarity = similarity[:, :long_mem_size]
# get the usage with the first group
# the first group always have all the keys valid
affinity, usage = do_softmax(
torch.cat([long_mem_similarity[:, -self.long_mem.get_v_size(0):], work_mem_similarity], 1),
top_k=self.top_k, inplace=True, return_usage=True)
affinity = [affinity]
# compute affinity group by group as later groups only have a subset of keys
for gi in range(1, num_groups):
if gi < self.long_mem.num_groups:
# merge working and lt similarities before softmax
affinity_one_group = do_softmax(
torch.cat([long_mem_similarity[:, -self.long_mem.get_v_size(gi):],
work_mem_similarity[:, -self.work_mem.get_v_size(gi):]], 1),
top_k=self.top_k, inplace=True)
else:
# no long-term memory for this group
affinity_one_group = do_softmax(work_mem_similarity[:, -self.work_mem.get_v_size(gi):],
top_k=self.top_k, inplace=(gi==num_groups-1))
affinity.append(affinity_one_group)
all_memory_value = []
for gi, gv in enumerate(self.work_mem.value):
# merge the working and lt values before readout
if gi < self.long_mem.num_groups:
all_memory_value.append(torch.cat([self.long_mem.value[gi], self.work_mem.value[gi]], -1))
else:
all_memory_value.append(gv)
"""
Record memory usage for working and long-term memory
"""
# ignore the index return for long-term memory
work_usage = usage[:, long_mem_size:]
self.work_mem.update_usage(work_usage.flatten())
if self.enable_long_term_usage:
# ignore the index return for working memory
long_usage = usage[:, :long_mem_size]
self.long_mem.update_usage(long_usage.flatten())
else:
# No long-term memory
similarity = get_similarity(self.work_mem.key, self.work_mem.shrinkage, query_key, selection)
if self.enable_long_term:
affinity, usage = do_softmax(similarity, inplace=(num_groups==1),
top_k=self.top_k, return_usage=True)
# Record memory usage for working memory
self.work_mem.update_usage(usage.flatten())
else:
affinity = do_softmax(similarity, inplace=(num_groups==1),
top_k=self.top_k, return_usage=False)
affinity = [affinity]
# compute affinity group by group as later groups only have a subset of keys
for gi in range(1, num_groups):
affinity_one_group = do_softmax(similarity[:, -self.work_mem.get_v_size(gi):],
top_k=self.top_k, inplace=(gi==num_groups-1))
affinity.append(affinity_one_group)
all_memory_value = self.work_mem.value
# Shared affinity within each group
all_readout_mem = torch.cat([
self._readout(affinity[gi], gv)
for gi, gv in enumerate(all_memory_value)
], 0)
return all_readout_mem.view(all_readout_mem.shape[0], self.CV, h, w)
def add_memory(self, key, shrinkage, value, objects, selection=None):
# key: 1*C*H*W
# value: 1*num_objects*C*H*W
# objects contain a list of object indices
if self.H is None or self.reset_config:
self.reset_config = False
self.H, self.W = key.shape[-2:]
self.HW = self.H*self.W
if self.enable_long_term:
# convert from num. frames to num. nodes
self.min_work_elements = self.min_mt_frames*self.HW
self.max_work_elements = self.max_mt_frames*self.HW
# key: 1*C*N
# value: num_objects*C*N
key = key.flatten(start_dim=2)
shrinkage = shrinkage.flatten(start_dim=2)
value = value[0].flatten(start_dim=2)
self.CK = key.shape[1]
self.CV = value.shape[1]
if selection is not None:
if not self.enable_long_term:
warnings.warn('the selection factor is only needed in long-term mode', UserWarning)
selection = selection.flatten(start_dim=2)
self.work_mem.add(key, value, shrinkage, selection, objects)
# long-term memory cleanup
if self.enable_long_term:
# Do memory compressed if needed
if self.work_mem.size >= self.max_work_elements:
# print('remove memory')
# Remove obsolete features if needed
if self.long_mem.size >= (self.max_long_elements-self.num_prototypes):
self.long_mem.remove_obsolete_features(self.max_long_elements-self.num_prototypes)
self.compress_features()
def create_hidden_state(self, n, sample_key):
# n is the TOTAL number of objects
h, w = sample_key.shape[-2:]
if self.hidden is None:
self.hidden = torch.zeros((1, n, self.hidden_dim, h, w), device=sample_key.device)
elif self.hidden.shape[1] != n:
self.hidden = torch.cat([
self.hidden,
torch.zeros((1, n-self.hidden.shape[1], self.hidden_dim, h, w), device=sample_key.device)
], 1)
assert(self.hidden.shape[1] == n)
def set_hidden(self, hidden):
self.hidden = hidden
def get_hidden(self):
return self.hidden
def compress_features(self):
HW = self.HW
candidate_value = []
total_work_mem_size = self.work_mem.size
for gv in self.work_mem.value:
# Some object groups might be added later in the video
# So not all keys have values associated with all objects
# We need to keep track of the key->value validity
mem_size_in_this_group = gv.shape[-1]
if mem_size_in_this_group == total_work_mem_size:
# full LT
candidate_value.append(gv[:,:,HW:-self.min_work_elements+HW])
else:
# mem_size is smaller than total_work_mem_size, but at least HW
assert HW <= mem_size_in_this_group < total_work_mem_size
if mem_size_in_this_group > self.min_work_elements+HW:
# part of this object group still goes into LT
candidate_value.append(gv[:,:,HW:-self.min_work_elements+HW])
else:
# this object group cannot go to the LT at all
candidate_value.append(None)
# perform memory consolidation
prototype_key, prototype_value, prototype_shrinkage = self.consolidation(
*self.work_mem.get_all_sliced(HW, -self.min_work_elements+HW), candidate_value)
# remove consolidated working memory
self.work_mem.sieve_by_range(HW, -self.min_work_elements+HW, min_size=self.min_work_elements+HW)
# add to long-term memory
self.long_mem.add(prototype_key, prototype_value, prototype_shrinkage, selection=None, objects=None)
# print(f'long memory size: {self.long_mem.size}')
# print(f'work memory size: {self.work_mem.size}')
def consolidation(self, candidate_key, candidate_shrinkage, candidate_selection, usage, candidate_value):
# keys: 1*C*N
# values: num_objects*C*N
N = candidate_key.shape[-1]
# find the indices with max usage
_, max_usage_indices = torch.topk(usage, k=self.num_prototypes, dim=-1, sorted=True)
prototype_indices = max_usage_indices.flatten()
# Prototypes are invalid for out-of-bound groups
validity = [prototype_indices >= (N-gv.shape[2]) if gv is not None else None for gv in candidate_value]
prototype_key = candidate_key[:, :, prototype_indices]
prototype_selection = candidate_selection[:, :, prototype_indices] if candidate_selection is not None else None
"""
Potentiation step
"""
similarity = get_similarity(candidate_key, candidate_shrinkage, prototype_key, prototype_selection)
# convert similarity to affinity
# need to do it group by group since the softmax normalization would be different
affinity = [
do_softmax(similarity[:, -gv.shape[2]:, validity[gi]]) if gv is not None else None
for gi, gv in enumerate(candidate_value)
]
# some values can be have all False validity. Weed them out.
affinity = [
aff if aff is None or aff.shape[-1] > 0 else None for aff in affinity
]
# readout the values
prototype_value = [
self._readout(affinity[gi], gv) if affinity[gi] is not None else None
for gi, gv in enumerate(candidate_value)
]
# readout the shrinkage term
prototype_shrinkage = self._readout(affinity[0], candidate_shrinkage) if candidate_shrinkage is not None else None
return prototype_key, prototype_value, prototype_shrinkage