1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
|
"""
Train ViT-Mini with block-level DFA on CIFAR-10 and SAVE the final checkpoint
+ the random feedback Bs. The existing snapshot_evolution_vit.py and
vit_frozen_blocks_baseline.py scripts do not save model checkpoints, which
means the protocol cannot be applied to a trained ViT post-hoc.
Output:
results/vit_dfa_checkpoints/dfa_vit_s{seed}.pt — state_dict + Bs
Run:
CUDA_VISIBLE_DEVICES=2 python experiments/train_vit_dfa_save_checkpoint.py --seed 42 --epochs 60
"""
import sys, os, argparse
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms
import numpy as np
from models.vit_mini import ViTMini
def get_loaders(batch_size=128):
tv_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
])
tv = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
])
tr = torchvision.datasets.CIFAR10('./data', True, download=True, transform=tv_train)
te = torchvision.datasets.CIFAR10('./data', False, download=True, transform=tv)
return (
DataLoader(tr, batch_size=batch_size, shuffle=True, num_workers=2),
DataLoader(te, batch_size=batch_size, shuffle=False, num_workers=2),
)
def evaluate(model, loader, dev):
model.eval()
n = c = 0
with torch.no_grad():
for x, y in loader:
x, y = x.to(dev), y.to(dev)
preds = model(x).argmax(-1)
c += (preds == y).sum().item()
n += x.size(0)
return c / n
def train_dfa_vit(model, train_loader, test_loader, dev, epochs, lr, wd):
d_model = model.d_hidden
L = model.num_blocks
C = 10
Bs = [torch.randn(d_model, C, device=dev) / np.sqrt(C) for _ in range(L)]
block_opts = [optim.AdamW(b.parameters(), lr=lr, weight_decay=wd) for b in model.blocks]
embed_opt = optim.AdamW(
list(model.patch_embed.parameters()) + [model.cls_token, model.pos_embed],
lr=lr, weight_decay=wd)
head_opt = optim.AdamW(
list(model.out_head.parameters()) + list(model.out_ln.parameters()),
lr=lr, weight_decay=wd)
scheds = [optim.lr_scheduler.CosineAnnealingLR(o, T_max=epochs) for o in block_opts] + [
optim.lr_scheduler.CosineAnnealingLR(embed_opt, T_max=epochs),
optim.lr_scheduler.CosineAnnealingLR(head_opt, T_max=epochs),
]
for ep in range(1, epochs + 1):
model.train()
for x, y in train_loader:
x, y = x.to(dev), y.to(dev)
batch = x.size(0)
with torch.no_grad():
logits, hiddens = model(x, return_hidden=True)
e_T = logits.softmax(-1); e_T[torch.arange(batch), y] -= 1
hL_det = hiddens[-1].detach()
h_cls = model.out_ln(hL_det[:, 0])
head_opt.zero_grad()
F.cross_entropy(model.out_head(h_cls), y).backward()
head_opt.step()
for l in range(L):
h_l = hiddens[l].detach()
a_dfa = (e_T @ Bs[l].T).detach()
a_dfa_b = a_dfa.unsqueeze(1).expand_as(h_l)
rms = (a_dfa_b ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
a_norm = a_dfa_b / rms
f_l = model.blocks[l](h_l)
local = (f_l * a_norm).sum(dim=-1).mean()
block_opts[l].zero_grad()
local.backward()
torch.nn.utils.clip_grad_norm_(model.blocks[l].parameters(), 1.0)
block_opts[l].step()
a0 = (e_T @ Bs[0].T).detach()
rms0 = (a0 ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
h0 = model.embed(x)
a0_b = a0.unsqueeze(1).expand_as(h0)
embed_loss = (h0 * (a0_b / rms0.unsqueeze(1))).sum(dim=-1).mean()
embed_opt.zero_grad()
embed_loss.backward()
embed_opt.step()
for s in scheds: s.step()
if ep % 10 == 0 or ep == 1 or ep == epochs:
acc = evaluate(model, test_loader, dev)
print(f" ep {ep}: test_acc={acc:.4f}", flush=True)
return Bs
def main():
p = argparse.ArgumentParser()
p.add_argument('--seed', type=int, default=42)
p.add_argument('--epochs', type=int, default=60)
p.add_argument('--lr', type=float, default=1e-3)
p.add_argument('--wd', type=float, default=0.05)
p.add_argument('--output_dir', type=str, default='results/vit_dfa_checkpoints')
args = p.parse_args()
os.makedirs(args.output_dir, exist_ok=True)
dev = torch.device('cuda:0')
print(f"Train ViT-Mini DFA: seed={args.seed} epochs={args.epochs}", flush=True)
train_loader, test_loader = get_loaders(batch_size=128)
torch.manual_seed(args.seed); np.random.seed(args.seed); torch.cuda.manual_seed_all(args.seed)
m = ViTMini(num_blocks=4, d_model=128, n_heads=4).to(dev)
Bs = train_dfa_vit(m, train_loader, test_loader, dev, args.epochs, args.lr, args.wd)
final_acc = evaluate(m, test_loader, dev)
print(f"FINAL test acc: {final_acc:.4f}", flush=True)
out_path = os.path.join(args.output_dir, f"dfa_vit_s{args.seed}.pt")
torch.save({
"state_dict": m.state_dict(),
"Bs": [b.cpu() for b in Bs],
"config": vars(args),
"test_acc": final_acc,
}, out_path)
print(f"Saved {out_path}")
if __name__ == "__main__":
main()
|
