"""DINOv3 pilot: mutable output adapters for all four original tasks. Frozen DINOv3/DPT; preserves original logits at initialization. Partial labels are explicit: missing supervision is never silently converted to a negative. """ import os os.environ.update(TF_CPP_MIN_LOG_LEVEL='2',TF_NUM_INTRAOP_THREADS='6',TF_NUM_INTEROP_THREADS='2') import json,hashlib from pathlib import Path import tensorflow as tf,numpy as np from merge_training_graph import merge,dynamic_resize_backbone root=Path('/workspace/vds-litert-batch2');key='fireviewer_dinov3_pointing_pilot_v1';out=root/'learning'/key;out.mkdir(exist_ok=True) S=448 names=['abstention_logits','point_logits','presence_logits','segmentation_logits'] shapes=[[1],[1,1,S,S],[1,2],[1,1,S,S]] features=[tf.TensorSpec(shape,tf.float32,name='features_'+name) for name,shape in zip(names,shapes)] class Head(tf.Module): def __init__(self): super().__init__() self.scale=tf.Variable(tf.ones([4]),name='scale');self.bias=tf.Variable(tf.zeros([4]),name='bias') self.pw=tf.Variable(tf.eye(2),name='presence_weight');self.pb=tf.Variable(tf.zeros([2]),name='presence_bias') self.parameters=[self.scale,self.bias,self.pw,self.pb] self.momentum=[tf.Variable(tf.zeros_like(v),name='momentum_'+v.name.split(':')[0]) for v in self.parameters] self.step=tf.Variable(0,dtype=tf.int64,name='step');self.state=self.parameters+self.momentum+[self.step] self.state_names=[v.name.split(':')[0] for v in self.state] self.infer=tf.function(self._infer,input_signature=features) self.train=tf.function(self._train,input_signature=features+[ tf.TensorSpec([1,3,S,S],tf.float32,name='y'), tf.TensorSpec([1,2],tf.float32,name='presence'), tf.TensorSpec([1],tf.float32,name='abstention'), tf.TensorSpec([4],tf.float32,name='supervision'), tf.TensorSpec([],tf.float32,name='learning_rate')]) def _infer(self,*f): return {n:(v@self.pw+self.pb if n=='presence_logits' else v*self.scale[i]+self.bias[i]) for i,(n,v) in enumerate(zip(names,f))} def _train(self,*args): y,presence,abstention,supervision,rate=args[-5:] with tf.GradientTape() as tape: outputs=self._infer(*args[:-5]);valid=y[:,2:3] spatial=[] for name,target in [('segmentation_logits',y[:,0:1]),('point_logits',y[:,1:2])]: loss=tf.nn.sigmoid_cross_entropy_with_logits(labels=target,logits=outputs[name]) spatial.append(tf.reduce_sum(loss*valid)/tf.maximum(tf.reduce_sum(valid),1.)) point,seg=spatial[1],spatial[0] abst=tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=abstention,logits=outputs['abstention_logits'])) pres=tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=presence,logits=outputs['presence_logits'])) loss=tf.reduce_sum(tf.stack([seg,point,abst,pres])*supervision*tf.constant([1.,.5,.25,.25])) gradients=tape.gradient(loss,self.parameters);gradients,_=tf.clip_by_global_norm(gradients,10.) for v,m,g in zip(self.parameters,self.momentum,gradients):m.assign(.9*m+g);v.assign_sub(rate*m) self.step.assign_add(1);return {'loss':loss,'step':self.step.read_value()} @tf.function(input_signature=[tf.TensorSpec([],tf.string,name='checkpoint_path')]) def save(self,checkpoint_path): tf.raw_ops.Save(filename=checkpoint_path,tensor_names=self.state_names,data=[v.read_value() for v in self.state]);return {'step':self.step.read_value()} @tf.function(input_signature=[tf.TensorSpec([],tf.string,name='checkpoint_path')]) def restore(self,checkpoint_path): for n,v in zip(self.state_names,self.state):v.assign(tf.raw_ops.Restore(file_pattern=checkpoint_path,tensor_name=n,dt=v.dtype)) return {'step':self.step.read_value()} module=Head();signatures=['infer','train','save','restore'];saved=out/'head_saved_model' tf.saved_model.save(module,str(saved),signatures={n:getattr(module,n).get_concrete_function() for n in signatures}) converter=tf.lite.TFLiteConverter.from_saved_model(str(saved));converter.target_spec.supported_ops=[tf.lite.OpsSet.TFLITE_BUILTINS,tf.lite.OpsSet.SELECT_TF_OPS];converter.experimental_enable_resource_variables=True head=out/'head.tflite';head.write_bytes(converter.convert());base=root/'learning-backbones'/key/'model.tflite' dynamic=out/'dynamic_backbone.tflite';preprocessing=dynamic_resize_backbone(base,dynamic) artifact=out/'model.tflite';merge(dynamic,head,artifact,{'features_'+n:i for i,n in enumerate(names)},image_size=S) def load(): i=tf.lite.Interpreter(model_path=str(artifact),num_threads=6,experimental_op_resolver_type=tf.lite.experimental.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES) return i,{n:i.get_signature_runner(n) for n in signatures} i,s=load();x=np.random.default_rng(94).normal(size=(1,3,S,S)).astype(np.float32);before=s['infer'](x=x) reference=tf.lite.Interpreter(model_path=str(base),num_threads=6,experimental_op_resolver_type=tf.lite.experimental.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES);reference.allocate_tensors();reference.set_tensor(reference.get_input_details()[0]['index'],x);reference.invoke() for name,d in zip(names,reference.get_output_details()):np.testing.assert_allclose(before[name],reference.get_tensor(d['index']),atol=2e-3,rtol=2e-3) y=np.zeros((1,3,S,S),np.float32);y[:,2]=1;y[:,0,150:300,140:280]=1;y[:,1,210:230,200:220]=1 targets={'y':y,'presence':np.array([[1,0]],np.float32),'abstention':np.array([0],np.float32),'supervision':np.ones(4,np.float32)} losses=[float(s['train'](x=x,**targets,learning_rate=np.float32(.01))['loss']) for _ in range(6)] assert np.isfinite(losses).all() and losses[-1]