diff --git a/pilot/analyze.py b/pilot/analyze.py index 0955e30..04854b5 100644 --- a/pilot/analyze.py +++ b/pilot/analyze.py @@ -58,7 +58,10 @@ def sample_object(obj, tensors, sample_size=2048): def exact_pair(a, b, tensors, deadline): if a['kind'] != 'tensor': - return metrics(read_object(a, tensors), read_object(b, tensors)) + result = metrics(read_object(a, tensors), read_object(b, tensors)) + tensors[a['tensor']].evict() + tensors[b['tensor']].evict() + return result ta, tb = tensors[a['tensor']], tensors[b['tensor']] xx = yy = xy = delta = sx = sy = 0. for start in range(0, ta.numel, 262144): @@ -227,7 +230,10 @@ def structural_probe(a, b, tensors): def analyze_mode(mode, objects, tensors, root, deadline, cap, seed, baseline_cap=100): print(f'{mode}: sampling {len(objects)} objects', flush=True) - features = [sample_object(o, tensors) for o in objects] + features = [] + for obj in objects: + features.append(sample_object(obj, tensors)) + tensors[obj['tensor']].evict() candidates, coverage = nearest_candidates(objects, features, cap) write_json(root / f'{mode}_objects.json', objects) pairs, nulls, structures = [], [], [] @@ -273,6 +279,8 @@ def analyze_mode(mode, objects, tensors, root, deadline, cap, seed, baseline_cap ndim = len(objects[a]['shape']) if ndim >= 2 and structural_counts[ndim] < 20: structures.append({'a': a, 'b': b, 'modes': structural_probe(objects[a], objects[b], tensors)}) + tensors[objects[a]['tensor']].evict() + tensors[objects[b]['tensor']].evict() structural_counts[ndim] += 1 if (i + 1) % 100 == 0: write_csv(root / f'{mode}_pairs.csv', pairs)