File size: 9,465 Bytes
e809d15 | 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 173 174 175 176 177 178 179 180 | """Master: masked IMU reconstruction and supervised local/global odometry."""
import math
import torch
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from .lab_bigru.model import WindowEncoder
class TCNBlock(nn.Module):
def __init__(self, hidden, dilation):
super().__init__()
self.norm=nn.LayerNorm(hidden)
self.conv=nn.Conv1d(hidden,hidden,3,padding=dilation,dilation=dilation)
self.proj=nn.Linear(hidden,hidden)
nn.init.zeros_(self.proj.weight);nn.init.zeros_(self.proj.bias)
def forward(self,x,valid):
z=self.norm(x)*valid[:,:,None]
z=self.conv(z.transpose(1,2)).transpose(1,2)
return (x+self.proj(torch.nn.functional.gelu(z)))*valid[:,:,None]
class Master(nn.Module):
def __init__(self,mean,std,hidden=256):
super().__init__()
self.hidden=hidden
self.encoder=WindowEncoder(mean,std)
self.project=nn.Sequential(nn.Linear(384,hidden),nn.LayerNorm(hidden),nn.GELU(),nn.Dropout(.1))
self.tcn=nn.ModuleList([TCNBlock(hidden,d) for d in (1,2,4)])
self.temporal=nn.GRU(hidden,hidden,2,batch_first=True,bidirectional=True,dropout=.1)
self.norm=nn.LayerNorm(hidden*2)
layer=nn.TransformerEncoderLayer(hidden,4,hidden*4,.1,batch_first=True,norm_first=True,activation='gelu')
self.global_net=nn.TransformerEncoder(layer,2,enable_nested_tensor=False)
# Independent initial layer draws rather than cloned identical weights.
for layer in self.global_net.layers:
for module in layer.modules():
if isinstance(module,nn.Linear):module.reset_parameters()
layer.self_attn._reset_parameters()
self.fuse=nn.Linear(hidden,hidden*2)
nn.init.zeros_(self.fuse.weight);nn.init.zeros_(self.fuse.bias)
self.direction=nn.Sequential(nn.Linear(hidden*2,hidden),nn.GELU(),nn.Dropout(.1),nn.Linear(hidden,3))
self.speed=nn.Sequential(nn.Linear(hidden*2,hidden),nn.GELU(),nn.Dropout(.1),nn.Linear(hidden,1))
self.residual=nn.Sequential(nn.Linear(hidden*2,hidden//2),nn.GELU(),nn.Linear(hidden//2,3))
nn.init.zeros_(self.residual[-1].weight);nn.init.zeros_(self.residual[-1].bias)
self.mask_embedding=nn.Parameter(torch.zeros(hidden))
self.reconstruct=nn.Sequential(nn.Linear(hidden*2,hidden),nn.GELU(),nn.Linear(hidden,240))
def forward(self,raw,lengths,starts,mask=None):
b,t=raw.shape[:2];device=raw.device
valid=torch.arange(t,device=device)[None]<lengths.to(device)[:,None]
if mask is not None:
raw=torch.where(mask[:,:,None,None],self.encoder.input_mean[None,None,:,None],raw)
raw=torch.where(valid[:,:,None,None],raw, self.encoder.input_mean[None,None,:,None])
flat=raw.reshape(-1,12,20)
# Encoding chunks controls convolution workspace, without changing sample statistics.
enc=torch.cat([self.encoder(z) for z in flat.split(256)]).reshape(b,t,384)
x=self.project(enc)
if mask is not None:x=x+mask[:,:,None]*self.mask_embedding
x=x*valid[:,:,None]
for block in self.tcn:x=block(x,valid)
packed=pack_padded_sequence(x,lengths.cpu(),batch_first=True,enforce_sorted=False)
out,_=self.temporal(packed)
h,_=pad_packed_sequence(out,batch_first=True,total_length=t)
h=self.norm(h)
group=(torch.arange(t,device=device)[None]+starts.to(device)[:,None]%10)//10
groups=(t+18)//10
pooled=x.new_zeros(b,groups,self.hidden)
pooled.scatter_add_(1,group[:,:,None].expand(-1,-1,self.hidden),x*valid[:,:,None])
counts=x.new_zeros(b,groups)
counts.scatter_add_(1,group,valid.to(x.dtype))
pooled=pooled/counts.clamp_min(1)[:,:,None]
pos=torch.arange(groups,device=device,dtype=torch.float32)[:,None]
freq=torch.exp(torch.arange(0,self.hidden,2,device=device)*(-math.log(10000)/self.hidden))
pe=torch.zeros(groups,self.hidden,device=device)
pe[:,0::2]=torch.sin(pos*freq);pe[:,1::2]=torch.cos(pos*freq)
g=self.global_net(pooled+pe.to(pooled.dtype),src_key_padding_mask=counts==0)
g=g.gather(1,group[:,:,None].expand(-1,-1,self.hidden))
h=h+self.fuse(g)
if mask is not None:return self.reconstruct(h).reshape(b,t,12,20)
d=self.direction(h).float()
v=torch.nn.functional.normalize(d,dim=-1,eps=1e-6)*torch.nn.functional.softplus(self.speed(h).float())
return (v+self.residual(h).float())*valid[:,:,None]
def block_mask(lengths,time,generator):
result=torch.zeros(len(lengths),time,dtype=torch.bool)
for b,n in enumerate(lengths.tolist()):
if n<2:raise ValueError('need masked and visible tokens')
count=max(1,min(n-1,round(n*.2)))
while int(result[b].sum())<count:
start=int(torch.randint(n,(1,),generator=generator))
width=int(torch.randint(5,21,(1,),generator=generator))
for i in range(start,min(n,start+width)):
if int(result[b].sum())==count:break
result[b,i]=True
return result
def pretrain_loss(pred,raw,mask,model):
target=(raw.float()-model.encoder.input_mean[None,None,:,None])/model.encoder.input_std[None,None,:,None]
error=(pred.float()-target).square().mean((2,3))
counts=mask.sum(1)
if (counts==0).any():raise ValueError('empty reconstruction mask')
return ((error*mask).sum(1)/counts).mean()
def supervised_loss(pred,batch,ramp):
target=batch['target'];valid=torch.arange(pred.shape[1],device=pred.device)[None]<batch['lengths'].to(pred.device)[:,None]
speed=torch.linalg.vector_norm(pred,dim=-1);truth=torch.linalg.vector_norm(target,dim=-1)
err=torch.linalg.vector_norm(pred-target,dim=-1)
mean=lambda x:((x*valid).sum(1)/valid.sum(1)).mean()
dense=mean(err*torch.where(speed<truth,2.,1.))
mag=mean((speed-truth).abs())
ones=[];disp={k:[] for k in (2,5,10)}
for b,(n,start) in enumerate(zip(batch['lengths'].tolist(),batch['starts'].tolist())):
offset=(-start)%10;count=(n-offset)//10
if count<1:continue
delta=(pred[b,offset:offset+count*10]-target[b,offset:offset+count*10]).reshape(count,10,3).mean(1)
ones.append(torch.linalg.vector_norm(delta,dim=-1).mean())
rotation=batch['rotation'][b,offset+9:offset+count*10:10]
world=torch.einsum('nij,nj->ni',rotation.float(),delta.float())
cumulative=torch.cat([world.new_zeros(1,3),world.cumsum(0)])
for k in disp:
if count>=k:disp[k].append(torch.linalg.vector_norm(cumulative[k:]-cumulative[:-k],dim=-1).mean()/k)
zero=pred.sum()*0
one=torch.stack(ones).mean() if ones else zero
ds=[torch.stack(v).mean() for v in disp.values() if v]
displacement=torch.stack(ds).mean() if ds else zero
parts=dict(dense=dense,speed=mag,one_second=one,displacement=displacement)
return dense+ramp*(.25*mag+.25*one+.1*displacement),parts
def transfer(model,state):
excluded=('direction.','speed.','residual.','reconstruct.','mask_embedding')
own=model.state_dict(); keys={k for k in own if not k.startswith(excluded)}
source={k for k in state if not k.startswith(excluded)}
if keys!=source:raise ValueError('backbone key mismatch')
for k in keys:
if own[k].shape!=state[k].shape or own[k].dtype!=state[k].dtype:raise ValueError('backbone tensor mismatch '+k)
own[k]=state[k]
model.load_state_dict(own,strict=True)
def window_starts(n,window=640,stride=320):
if n<1:raise ValueError('empty trajectory')
if n<=window:return [0]
return sorted(set(list(range(0,n-window+1,stride))+[n-window]))
def validate_parent(payload,seed,regime,inputs,steps,smoke):
c=payload['contract']
if (c['seed']!=seed or c['regime']!=regime or c['phase']!='pretrain'
or c['historical_weights'] is not False or c['inputs']!=inputs
or (not smoke and (payload['step']!=steps or c['smoke']))):
raise ValueError('pretraining parent identity/data/step mismatch')
def load_f3_weights(model,state):
"""Strict mapping of every frozen F3 tensor; fresh residual paths stay zero."""
mapping={'encoder.':'encoder.','input_projection.':'project.','temporal.':'temporal.',
'output_norm.':'norm.','output.':'direction.','mag_head.':'speed.'}
own=model.state_dict();translated={}
for name,value in state.items():
prefix=next((p for p in mapping if name.startswith(p)),None)
if prefix is None:raise ValueError('unexpected F3 parameter '+name)
target=mapping[prefix]+name[len(prefix):]
if target not in own or value.shape!=own[target].shape or value.dtype!=own[target].dtype:
raise ValueError('F3 tensor mismatch '+name)
translated[target]=value
required={k for k in own if k.startswith(tuple(mapping.values()))}
if set(translated)!=required:raise ValueError('incomplete F3 checkpoint')
own.update(translated);model.load_state_dict(own,strict=True)
def validate_continuation(payload,seed,regime,inputs):
c=payload['contract']
if (c['seed']!=seed or c['regime']!=regime or c['phase']!='finetune'
or c['smoke'] or c['inputs']!=inputs or c['config']['name']!='f3_sd_warm'
or payload['step']!=18000 or c['config']['finetune_steps']!=18000):
raise ValueError('continuation requires matching completed Warm Master')
|