| |
| |
| from pathlib import Path |
| import sys,hashlib |
| import onnx,numpy as np |
| from onnx import helper as h,numpy_helper as nh |
| def export(base,path): |
| p=[1,1,.2,0,0,0,0,-.075,1] |
| p=np.array(p,np.float32);m=onnx.load(base);inp=m.graph.input[0].name;output=m.graph.output[0].name;nodes=[];init=[] |
| def c(name,value): |
| init.append(nh.from_array(np.array(value),name));return name |
| def n(op,inputs,out,**kw):nodes.append(h.make_node(op,inputs,[out],**kw));return out |
| def gather(i):return n('Gather',[inp,c('idx'+str(i),np.array([i],np.int64))],'slot'+str(i),axis=1) |
| vx=gather(48);yaw=gather(50);gx=gather(3);gz=gather(2) |
| green=n('Cast',[n('Not',[n('Equal',[vx,c('zero',np.float32(0))],'eq')],'neq')],'green',to=1) |
| y2=n('Mul',[yaw,yaw],'y2');den=n('Add',[c('one',np.float32(1)),n('Mul',[c('k',p[2]),y2],'ky2')],'den') |
| eff=n('Div',[n('Mul',[vx,c('vxscale',p[0])],'vxs'),den],'eff');vnew=n('Add',[vx,n('Sub',[eff,vx],'diff')],'vnew') |
| ynew=n('Mul',[yaw,c('yscale',p[1])],'ynew');gxnew=n('Add',[gx,n('Mul',[green,c('gxbias',p[3])],'gxd')],'gxnew');gznew=n('Mul',[gz,c('gzscale',p[8])],'gznew') |
| obs=inp |
| for i,val in [(48,vnew),(50,ynew),(3,gxnew),(2,gznew)]: |
| obs=n('ScatterElements',[obs,c('scatter'+str(i),np.array([[i]],np.int64)),val],'obs'+str(i),axis=1) |
| old=list(m.graph.node) |
| for node in old: |
| for i,v in enumerate(node.input): |
| if v==inp:node.input[i]=obs |
| for i,v in enumerate(node.output): |
| if v==output:node.output[i]='base_output' |
| bias=np.zeros((1,14),np.float32);bias[0,[2,3,4,5,6,11,12,13]]=[p[4],p[5],p[6],p[7],p[7],-p[4],-p[5],-p[6]] |
| delta=n('Mul',[green,c('bias',bias)],'act_delta') |
| |
| final=h.make_node('Add',['base_output',delta],[output]);del m.graph.node[:];m.graph.node.extend(nodes+old+[final]);m.graph.initializer.extend(init) |
| onnx.checker.check_model(m);onnx.save(m,path) |
|
|
| if __name__=='__main__': |
| if len(sys.argv)!=3:raise SystemExit('Usage: python export_selected.py UPSTREAM_alpha_walking.onnx OUTPUT.onnx') |
| base=Path(sys.argv[1]);out=Path(sys.argv[2]) |
| assert hashlib.sha256(base.read_bytes()).hexdigest()=='e36332d383997d51401897734cd3e79cf5038406feddb18b4d57ecfb141daa6c' |
| export(base,out) |
| assert hashlib.sha256(out.read_bytes()).hexdigest()=='e9300f63f8374029ea21c732f78f27a84b54926baabe0d240d13cd27a669f44a' |
| print('Exact selected policy bytes reproduced') |
|
|
