Spaces:
Runtime error
Runtime error
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 |