1import json
 2import random
 3import time
 4import numpy as np
 5import torch
 6from torch import nn
 7
 8SEED = 19
 9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
10try:
11    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
12except Exception:
13    device = torch.device('cpu')
14
15L, D, H, N = 30, 4, 8, 2400
16rng = np.random.default_rng(SEED)
17x = rng.normal(size=(N, L, D)).astype('float32')
18y = (x.sum(axis=(1, 2)) > 0).astype('int64')
19xt = torch.tensor(x); yt = torch.tensor(y)
20
21class StandardRNN(nn.Module):
22    def __init__(self):
23        super().__init__()
24        self.inp = nn.Linear(D, H)
25        self.rec = nn.Linear(H, H)
26        self.out = nn.Linear(H, 2)
27    def forward(self, x):
28        h = torch.zeros(x.size(0), H, device=x.device)
29        for t in range(L):
30            h = torch.tanh(self.inp(x[:, t]) + self.rec(h))
31        return self.out(h)
32
33class ContractingState(nn.Module):
34    # A=rI, with a bounded driver and learned forcing phi(x_t,u_t).
35    def __init__(self, r):
36        super().__init__()
37        self.r = r
38        self.wx = nn.Linear(D, 1)
39        self.force = nn.Sequential(nn.Linear(D + 1, 16), nn.Tanh(), nn.Linear(16, H))
40        self.out = nn.Linear(H, 2)
41    def forward(self, x):
42        h = torch.zeros(x.size(0), H, device=x.device)
43        for t in range(L):
44            driver = torch.sigmoid(self.wx(x[:, t]))
45            h = self.r * h + self.force(torch.cat([driver, x[:, t]], dim=1))
46        return self.out(h)
47
48def run(model, epochs=18):
49    model = model.to(device)
50    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
51    lossfn = nn.CrossEntropyLoss()
52    tic = time.time()
53    for _ in range(epochs):
54        perm = torch.randperm(1800, device='cpu')
55        for start in range(0, 1800, 64):
56            ind = perm[start:start+64]
57            loss = lossfn(model(xt[ind].to(device)), yt[ind].to(device))
58            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
59    with torch.no_grad():
60        pred = model(xt[1800:].to(device)).argmax(1).cpu()
61        acc = float((pred == yt[1800:]).float().mean())
62        test_loss = float(lossfn(model(xt[1800:].to(device)), yt[1800:].to(device)))
63    return {'test_accuracy': acc, 'test_loss': test_loss, 'seconds': time.time()-tic,
64            'parameters': sum(p.numel() for p in model.parameters())}
65
66def main():
67    global device
68    results = {'seed': SEED, 'device': str(device), 'task': 'sign of sequence-wide sum'}
69    try:
70        results['baseline'] = run(StandardRNN())
71        results['idea_r095'] = run(ContractingState(.95))
72    except Exception as e:
73        # CUDA allocation/runtime errors are handled by rerunning on CPU.
74        if device.type == 'cuda':
75            device = torch.device('cpu')
76            results['device_fallback'] = str(e)
77            results['baseline'] = run(StandardRNN())
78            results['idea_r095'] = run(ContractingState(.95))
79        else:
80            raise
81    with open('mini_results.json', 'w') as f: json.dump(results, f, indent=2)
82    print(json.dumps(results, indent=2))
83
84if __name__ == '__main__': main()