Use separable SPPF pooling in PyTorch, preserving square pooling for export - #26082

import copy
import json
import platform
import statistics
import types
import torch
from torch.utils.benchmark import Timer
from ultralytics import YOLO
from ultralytics.nn.modules import SPPF

torch.manual_seed(0)
torch.set_num_threads(8)
print(json.dumps({'torch': torch.__version__, 'cuda': torch.version.cuda, 'python':platform.python_version(), 'gpu':torch.cuda.get_device_name(0)}),flush=True)
results=[]
def original_forward(self,x):
    y=[self.cv1(x)]
    y.extend(self.m(y[-1]) for _ in range(getattr(self,"n",3)))
    y=self.cv2(torch.cat(y,1))
    return y+x if getattr(self,"add",False) else y

def bench(label, shape, dtype, device, threads, old, new, training=False):
    def run(m):
        if training:
            m.zero_grad(set_to_none=True)
            m(x).sum().backward()
            x.grad=None
        else:
            with torch.inference_mode():
                m(x)
    x=torch.randn(shape,device=device,dtype=dtype,requires_grad=training)
    funcs=[lambda:run(old),lambda:run(new)]
    for f in funcs:
        for _ in range(10): f()
    samples=[[],[]]
    for rep in range(3):
        for i in ([0,1] if rep%2==0 else [1,0]):
            samples[i].append(Timer(stmt='f()',globals={'f':funcs[i]},num_threads=threads).blocked_autorange(min_run_time=0.3).median*1000)
    a,b=map(statistics.median,samples)
    row=dict(label=label,shape=list(shape),dtype=str(dtype).split('.')[-1],device=device,threads=threads,mode='forward+backward' if training else 'inference',old_ms=a,new_ms=b,speedup=a/b,samples_ms=samples)
    results.append(row)
    print(json.dumps(row),flush=True)
    with open('/tmp/sppf-results.json','w') as f:json.dump(results,f,indent=2)

class Pool(torch.nn.Module):
    def __init__(self,separable):
        super().__init__();self.separable=separable
    def forward(self,x):
        for _ in range(3):
            x=(torch.nn.functional.max_pool2d(torch.nn.functional.max_pool2d(x,(1,5),1,(0,2)),(5,1),1,(2,0)) if self.separable else torch.nn.functional.max_pool2d(x,5,1,2))
        return x

for device,dtype,threads in [('cpu',torch.float32,1),('cpu',torch.float32,8),('cuda',torch.float32,8),('cuda',torch.float16,8)]:
    for shape in [(1,256,20,20),(8,256,20,20),(1,256,40,40)]:
        for training in [False,True]:
            old=Pool(False).to(device);new=Pool(True).to(device)
            bench('3 pooling iterations',shape,dtype,device,threads,old,new,training)
            old=SPPF(shape[1],shape[1]).to(device=device,dtype=dtype);old.export=True
            new=copy.deepcopy(old);new.export=False
            old.forward=types.MethodType(original_forward,old)
            old.train(training);new.train(training)
            bench('SPPF',shape,dtype,device,threads,old,new,training)
    for name in ['yolo26n','yolo26s']:
        old=YOLO(name+'.pt').model.fuse(verbose=False).eval().to(device=device,dtype=dtype)
        for m in old.modules():
            if isinstance(m,SPPF):m.export=True
        new=copy.deepcopy(old)
        for m in old.modules():
            if isinstance(m,SPPF):m.forward=types.MethodType(original_forward,m)
        for m in new.modules():
            if isinstance(m,SPPF):m.export=False
        for batch in [1,8]:
            bench(name,(batch,3,640,640),dtype,device,threads,old,new)
        del old,new

import copy
import json
import torch
import onnx
import onnxruntime as ort
import numpy as np
from ultralytics import YOLO
from ultralytics.nn.modules import SPPF

torch.set_num_threads(8)
torch.manual_seed(42)
count=0
for device,dtype in [('cpu',torch.float32),('cuda',torch.float32),('cuda',torch.float16)]:
    for k in [3,5,7]:
        for shape in [(2,8,1,1),(2,8,3,7),(2,8,20,20)]:
            for tied in [False,True]:
                x=(torch.randint(-2,3,shape,device=device).to(dtype) if tied else torch.randn(shape,device=device,dtype=dtype)).requires_grad_()
                z=x.detach().clone().requires_grad_()
                a=torch.nn.functional.max_pool2d(x,k,1,k//2)
                b=torch.nn.functional.max_pool2d(torch.nn.functional.max_pool2d(z,(1,k),1,(0,k//2)),(k,1),1,(k//2,0))
                torch.testing.assert_close(a,b,rtol=0,atol=0)
                grad=torch.ones_like(a)
                a.backward(grad);b.backward(grad)
                torch.testing.assert_close(x.grad,z.grad,rtol=0,atol=0)
                count+=1
    for k,n,shortcut in [(3,1,False),(5,3,False),(7,4,True)]:
        a=SPPF(16,16,k,n,shortcut).to(device=device,dtype=dtype).train();a.export=True
        b=copy.deepcopy(a);b.export=False
        x=torch.randn(2,16,11,13,device=device,dtype=dtype,requires_grad=True)
        z=x.detach().clone().requires_grad_()
        out=a(x);other=b(z)
        torch.testing.assert_close(out,other,rtol=0,atol=0)
        out.float().square().mean().backward();other.float().square().mean().backward()
        tolerance=dict(rtol=0.02,atol=0.0001) if dtype==torch.float16 else dict(rtol=0.0001,atol=1e-7)
        torch.testing.assert_close(x.grad,z.grad,**tolerance)
        for p,q in zip(a.parameters(),b.parameters()):torch.testing.assert_close(p.grad,q.grad,**tolerance)
        count+=1
print('POOL_AND_SPPF_PARITY_PASS',count,flush=True)
for name in ['yolo26n.pt','yolo11n.pt']:
    model=YOLO(name)
    original=copy.deepcopy(model.model).eval()
    for m in original.modules():
        if isinstance(m,SPPF):m.export=True
    model.model.eval()
    x=torch.randn(1,3,160,160)
    with torch.inference_mode():
        torch.testing.assert_close(original(x),model.model(x),rtol=0,atol=0)
    model.predict('ultralytics/assets/bus.jpg',imgsz=160,device=0,verbose=False,save=False)
    print('CHECKPOINT_AND_REAL_IMAGE_PASS',name,flush=True)
model=YOLO('yolo26n.pt')
p=model.export(format='onnx',imgsz=160,device='cpu',simplify=False,opset=17)
graph=onnx.load(p)
pools=[node for node in graph.graph.node if node.op_type=='MaxPool']
kernels=[[list(a.ints) for a in node.attribute if a.name=='kernel_shape'][0] for node in pools]
assert kernels==[[5,5]]*3,kernels
assert all(not m.export for m in model.model.modules() if isinstance(m,SPPF))
session=ort.InferenceSession(p,providers=['CPUExecutionProvider'],sess_options=(lambda o:(setattr(o,'intra_op_num_threads',8),o)[1])(ort.SessionOptions()))
x=np.random.default_rng(42).normal(size=(1,3,160,160)).astype(np.float32)
output=session.run(None,{'images':x})[0]
reference=copy.deepcopy(model.model).cpu().eval().fuse(verbose=False)
from ultralytics.nn.modules import Detect
for m in reference.modules():
    if isinstance(m,Detect):m.export=True;m.format='onnx';m.max_det=300
with torch.inference_mode():expected=reference(torch.from_numpy(x)).numpy()
np.testing.assert_allclose(output,expected,rtol=1e-3,atol=1e-3)
print('ONNX_3_SQUARE_POOLS_AND_RUNTIME_PARITY_PASS',flush=True)
p=model.export(format='torchscript',imgsz=160,device='cpu')
script=torch.jit.load(p)
graph=str(script.inlined_graph)
assert graph.count('aten::max_pool2d(')==3
with torch.inference_mode():
    np.testing.assert_allclose(script(torch.from_numpy(x)).numpy(),expected,rtol=1e-3,atol=1e-3)
print('TORCHSCRIPT_3_POOLS_AND_RUNTIME_PARITY_PASS',flush=True)