File size: 9,252 Bytes
38e555d 9a7563a 38e555d 9a7563a 38e555d 9a7563a 38e555d 9a7563a 38e555d | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """Fresh raw-PASTIS inference for every component of the 58.6346% ensemble.
Paths are explicit to support uploaded/local checkpoints. The published recipe
is fixed; this module never selects weights with test labels. Use the released
component hashes to reproduce the recorded system, not a newly trained replacement.
"""
import argparse
import gc
import json
import time
from pathlib import Path
import numpy as np
import torch
import config as C
from dataset import PASTISPatchDataset, pad_collate, val_patch_ids, test_patch_ids, crop_ij
from model import UTAE_ASPP, UTAEClassificationDual
from temporal_cnn import TemporalCNN
from training_feature_clusters import DepthCollector, probability
import foundation_transfer as F
import tessera_student as student
def load(path):
return torch.load(path, map_location='cpu', weights_only=False)
def spatial_network(path):
net = UTAEClassificationDual(UTAE_ASPP(n_channels=C.N_CHANNELS, d_model=C.D_MODEL,
n_heads=C.N_HEADS), d_model=C.D_MODEL,
num_classes=C.N_CLASSES)
net.load_state_dict(load(path)['model_state_dict'], strict=True)
return net.eval().requires_grad_(False)
def class_probabilities(logits):
b, _, h, w = logits.shape
result = torch.zeros((b, 20, h, w), device=logits.device)
result[:, 1:19] = logits.float().softmax(1)
return result
class CompleteEnsemble:
"""Sequential GPU residency keeps the two spatial networks from overlapping."""
def __init__(self, paths):
if not torch.cuda.is_available():
raise RuntimeError('Fresh complete-ensemble inference requires CUDA.')
self.paths = {k: Path(v) for k, v in paths.items()}
for name, path in self.paths.items():
if not path.is_file(): raise FileNotFoundError(f'{name}: {path}')
self.sam = spatial_network(self.paths['sam'])
self.structured = spatial_network(self.paths['structured'])
self.collector = DepthCollector(self.sam.utae)
self.cluster = load(self.paths['clusters'])
self.old = TemporalCNN(18).cuda().eval().requires_grad_(False)
self.old.load_state_dict(load(self.paths['old_cnn'])['state'], strict=True)
self.new = TemporalCNN(18).cuda().eval().requires_grad_(False)
self.new.load_state_dict(load(self.paths['new_cnn'])['state'], strict=True)
self.foundation = student.load_model(str(self.paths['tessera']), torch.device('cuda')).eval().requires_grad_(False)
self.original_gaussian = dict(np.load(self.paths['original_gaussian']))
self.optical_gaussian = dict(np.load(self.paths['optical_gaussian']))
self.mean, self.std = self.cluster['mean'].cuda(), self.cluster['std'].cuda()
self.cluster_state = ({k: v.cuda() for k, v in self.cluster['state'].items()}
if self.cluster['kind'] == 'linear' else self.cluster['state'].cuda())
@torch.inference_mode()
def predict(self, x, pos, days, split, return_features=False):
"""Return B×H×W×20 probabilities; labels are never an argument."""
self.sam.cuda(); self.collector.begin(x, pos)
with torch.autocast('cuda', dtype=torch.float16):
output = self.sam(x, pos)
old_logits = self.old(x, pos)
new_logits = self.new(x, pos)
features = self.collector.end()
semantic = output['refined_logits'].float().softmax(1)
cluster = probability(features, self.cluster['depth'], self.mean, self.std,
self.cluster['kind'], self.cluster_state)
b, _, h, w = semantic.shape
cluster = torch.from_numpy(cluster).cuda().reshape(b, h, w, 20).permute(0, 3, 1, 2)
# Preserve original nesting and GPU float32 arithmetic.
spatial_cluster = .5 * semantic + .5 * cluster
old_reference = .75 * spatial_cluster + .25 * class_probabilities(old_logits)
reference = .75 * old_reference + .25 * class_probabilities(new_logits)
reference = reference.permute(0, 2, 3, 1).cpu().numpy()
del output, semantic, cluster, spatial_cluster, old_reference, old_logits, new_logits
self.sam.cpu(); torch.cuda.empty_cache()
# The original Gaussian classifier reads first-block temporal mean/std.
p = self.original_gaussian
z = (features[:, 0].astype(np.float32) - p['mean']) / p['std']
scores = z @ p['weight'] + p['bias']
gaussian = np.zeros((len(z), 20), np.float32)
gaussian[:, 1:19] = torch.from_numpy(scores / float(p['temperature'])).softmax(1).numpy()
gaussian = gaussian.reshape(b, h, w, 20)
self.structured.cuda()
with torch.autocast('cuda', dtype=torch.float16):
structured = self.structured(x, pos)['refined_logits'].float().softmax(1)
structured = structured.permute(0, 2, 3, 1).cpu().numpy()
self.structured.cpu(); torch.cuda.empty_cache()
optical = []; optical_features = []
for i in range(b):
inputs = F.prepare(x[i:i+1], pos[i:i+1], days[i:i+1])
f = F.optical_features(self.foundation, inputs)
# Original validation came from float16 feature caches; test was fresh.
if split == 'val': f = f.astype(np.float16).astype(np.float32)
optical_features.append(f.reshape(h, w, -1))
optical.append(F.predict(f, self.optical_gaussian).reshape(h, w, 20))
optical = np.stack(optical)
incumbent = .65 * reference + .25 * gaussian + .1 * structured
final = .5 * incumbent + .5 * optical
result = {'reference': reference, 'gaussian': gaussian, 'structured': structured,
'optical': optical, 'incumbent': incumbent, 'final': final}
if return_features: result['optical_features'] = np.stack(optical_features)
return result
def close(self):
self.collector.close()
for name in ['sam', 'structured', 'old', 'new', 'foundation']:
getattr(self, name).cpu()
gc.collect(); torch.cuda.empty_cache()
def evaluate(paths, split='val', output=None, max_batches=0, save_probabilities=False):
"""A limited smoke run is explicitly marked incomplete and has no headline score."""
from torch.utils.data import DataLoader
assert split in ['val', 'test']
torch.set_num_threads(4); torch.manual_seed(3407)
torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True
ids = val_patch_ids() if split == 'val' else test_patch_ids()
dl = DataLoader(PASTISPatchDataset(ids), batch_size=2 if split == 'val' else 1,
shuffle=False, num_workers=2, pin_memory=True, collate_fn=pad_collate)
model = CompleteEnsemble(paths); matrices = {}; offset = 0; entries = 0; chunks = []
started = time.monotonic()
try:
for bi, ((x, pos, days), y) in enumerate(dl):
if max_batches and bi >= max_batches: break
x, pos = x.cuda(), pos.cuda(); entries += len(x)
for i, j in C.FT_CROP_IJ if split == 'val' else [(0, 0)]:
crop, target = crop_ij(x, y, i, j, 4 if split == 'val' else 1)
parts = model.predict(crop, pos, days, split)
for name, p in parts.items():
matrices.setdefault(name, np.zeros((20, 20), np.int64))
matrices[name] += F.confusion(target.numpy().ravel(), p.argmax(-1).ravel())
if save_probabilities: chunks.append(parts['final'].reshape(-1, 20))
offset += target.numel()
if bi % 10 == 0: print(split, bi, len(dl), 'seconds', time.monotonic() - started, flush=True)
finally:
model.close()
complete = entries == len(ids)
result = {'split': split, 'complete': complete, 'patch_entries': entries,
'expected_patch_entries': len(ids), 'seed': 3407, 'pixel_entries': offset,
'seconds': time.monotonic() - started,
'components': {k: F.metrics(v) for k, v in matrices.items()},
'paths': {k: str(v) for k, v in paths.items()},
'selection': 'published fixed recipe; no test tuning'}
if output:
output = Path(output); output.mkdir(parents=True, exist_ok=True)
(output / f'fresh_{split}_results.json').write_text(json.dumps(result, indent=2))
if save_probabilities: np.save(output / f'fresh_{split}_probabilities.npy', np.concatenate(chunks))
return result
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--paths', type=Path, required=True, help='JSON mapping eight named component paths')
parser.add_argument('--split', choices=['val', 'test'], default='val')
parser.add_argument('--output', type=Path, required=True)
parser.add_argument('--max_batches', type=int, default=0)
parser.add_argument('--save_probabilities', action='store_true')
args = parser.parse_args()
results = evaluate(json.loads(args.paths.read_text()), args.split, args.output,
args.max_batches, args.save_probabilities)
print('complete', results['complete'], 'mIoU', results['components']['final']['miou'], flush=True)
|