1import json, random, time
 2import numpy as np
 3import torch
 4from torch import nn
 5from third_order_optimizer import ThirdOrderLangevin
 6
 7SEED=123
 8
 9def seed(s):
10    random.seed(s); np.random.seed(s); torch.manual_seed(s)
11    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
12
13def make_data(n=2048):
14    g=np.random.default_rng(77)
15    x=g.normal(size=(n,2)).astype('float32')
16    y=((x[:,0]*x[:,1] + .25*x[:,0] - .15*x[:,1])>0).astype('int64')
17    return torch.from_numpy(x), torch.from_numpy(y)
18
19def run(kind, device):
20    seed(SEED); x,y=make_data(); x,y=x.to(device),y.to(device)
21    model=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,2)).to(device)
22    if kind=='sgd': opt=torch.optim.SGD(model.parameters(),lr=.08)
23    elif kind=='momentum': opt=torch.optim.SGD(model.parameters(),lr=.04,momentum=.9)
24    else: opt=ThirdOrderLangevin(model.parameters(),dt=.035,gamma=1.0,temperature=1e-4)
25    lossfn=nn.CrossEntropyLoss(); losses=[]; t0=time.perf_counter()
26    for step in range(300):
27        opt.zero_grad(set_to_none=True); loss=lossfn(model(x),y); loss.backward(); opt.step()
28        if step%50==0: losses.append(float(loss.detach().cpu()))
29    with torch.no_grad(): acc=float((model(x).argmax(1)==y).float().mean().cpu())
30    return {'loss_checkpoints':losses,'final_loss':float(loss.detach().cpu()),'train_accuracy':acc,'seconds':time.perf_counter()-t0}
31
32def main():
33    device='cuda' if torch.cuda.is_available() else 'cpu'
34    try:
35        out={k:run(k,device) for k in ['sgd','momentum','third_order']}
36    except Exception:
37        device='cpu'; out={k:run(k,device) for k in ['sgd','momentum','third_order']}
38    result={'device':device,'steps':300,'results':out}
39    with open('benchmark_results.json','w') as f: json.dump(result,f,indent=2)
40    print(json.dumps(result,indent=2))
41if __name__=='__main__': main()