Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| import warnings | |
| from .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 |