Skip to content

Commit cde1279

Browse files
strands-agentstrands-agent
committed
Merge PR facebookresearch#121: Add ST-A² attention
2 parents 01a7961 + 26f8a61 commit cde1279

16 files changed

Lines changed: 4934 additions & 9 deletions

‎app/vjepa/train.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,12 @@ def main(args, resume_preempt=False):
9696
use_silu = cfgs_model.get("use_silu", False)
9797
use_pred_silu = cfgs_model.get("use_pred_silu", False)
9898
wide_silu = cfgs_model.get("wide_silu", True)
99+
# -- ST-A² (Spatiotemporal Area Attention)
100+
use_area_attention = cfgs_model.get("use_area_attention", False)
101+
area_attention_layers = cfgs_model.get("area_attention_layers", None)
102+
area_spatial_splits = cfgs_model.get("area_spatial_splits", 2)
103+
area_temporal_splits = cfgs_model.get("area_temporal_splits", 2)
104+
area_residual_scale = cfgs_model.get("area_residual_scale", 1.0)
99105

100106
# -- DATA
101107
cfgs_data = args.get("data")
@@ -218,6 +224,11 @@ def main(args, resume_preempt=False):
218224
wide_silu=wide_silu,
219225
use_rope=use_rope,
220226
use_activation_checkpointing=use_activation_checkpointing,
227+
use_area_attention=use_area_attention,
228+
area_attention_layers=area_attention_layers,
229+
area_spatial_splits=area_spatial_splits,
230+
area_temporal_splits=area_temporal_splits,
231+
area_residual_scale=area_residual_scale,
221232
)
222233
target_encoder = copy.deepcopy(encoder)
223234

‎app/vjepa/utils.py‎

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -104,27 +104,36 @@ def load_checkpoint(
104104
epoch = checkpoint["epoch"]
105105

106106
# -- loading encoder
107+
# Use strict=False when annealing to allow loading baseline checkpoints
108+
# into area-attention models (RoPEAreaAttention has identical weight
109+
# structure to RoPEAttention, so all shared params load correctly).
107110
pretrained_dict = checkpoint["encoder"]
108-
msg = encoder.load_state_dict(pretrained_dict)
111+
msg = encoder.load_state_dict(pretrained_dict, strict=not is_anneal)
109112
logger.info(f"loaded pretrained encoder from epoch {epoch} with msg: {msg}")
110113

111114
# -- loading predictor
112115
pretrained_dict = checkpoint["predictor"]
113-
msg = predictor.load_state_dict(pretrained_dict)
116+
msg = predictor.load_state_dict(pretrained_dict, strict=not is_anneal)
114117
logger.info(f"loaded pretrained predictor from epoch {epoch} with msg: {msg}")
115118

116119
# -- loading target_encoder
117120
if target_encoder is not None:
118121
print(list(checkpoint.keys()))
119122
pretrained_dict = checkpoint["target_encoder"]
120-
msg = target_encoder.load_state_dict(pretrained_dict)
123+
msg = target_encoder.load_state_dict(pretrained_dict, strict=not is_anneal)
121124
logger.info(f"loaded pretrained target encoder from epoch {epoch} with msg: {msg}")
122125

123126
# -- loading optimizer
124-
opt.load_state_dict(checkpoint["opt"])
125-
if scaler is not None:
126-
scaler.load_state_dict(checkpoint["scaler"])
127-
logger.info(f"loaded optimizers from epoch {epoch}")
127+
# Skip optimizer/scaler restore when annealing from a different
128+
# architecture (e.g., baseline → area-attention) because the optimizer
129+
# state dict keys won't match the new parameter set.
130+
if is_anneal:
131+
logger.info("Annealing: skipping optimizer/scaler restore (fresh optimizer)")
132+
else:
133+
opt.load_state_dict(checkpoint["opt"])
134+
if scaler is not None:
135+
scaler.load_state_dict(checkpoint["scaler"])
136+
logger.info(f"loaded optimizers from epoch {epoch}")
128137
logger.info(f"read-path: {r_path}")
129138
del checkpoint
130139

@@ -158,6 +167,12 @@ def init_video_model(
158167
use_pred_silu=False,
159168
wide_silu=False,
160169
use_activation_checkpointing=False,
170+
# -- ST-A² params
171+
use_area_attention=False,
172+
area_attention_layers=None,
173+
area_spatial_splits=2,
174+
area_temporal_splits=2,
175+
area_residual_scale=1.0,
161176
):
162177
encoder = video_vit.__dict__[model_name](
163178
img_size=crop_size,
@@ -170,6 +185,11 @@ def init_video_model(
170185
wide_silu=wide_silu,
171186
use_activation_checkpointing=use_activation_checkpointing,
172187
use_rope=use_rope,
188+
use_area_attention=use_area_attention,
189+
area_attention_layers=area_attention_layers,
190+
area_spatial_splits=area_spatial_splits,
191+
area_temporal_splits=area_temporal_splits,
192+
area_residual_scale=area_residual_scale,
173193
)
174194
encoder = MultiSeqWrapper(encoder)
175195
predictor = vit_pred.__dict__["vit_predictor"](

‎configs/eval/vitl/k400-area-attn.yaml‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
# ST-A² (Spatiotemporal Area Attention) eval config for Kinetics-400
2+
# Based on k400.yaml with area attention enabled in the encoder.
3+
#
4+
# The encoder checkpoint must have been trained with matching area attention
5+
# settings (use_area_attention=true, layers 0-17, 2x2 splits).
6+
#
7+
# Usage:
8+
# python -m evals.main --fname configs/eval/vitl/k400-area-attn.yaml \
9+
# --devices cuda:0 cuda:1 cuda:2 cuda:3 cuda:4 cuda:5 cuda:6 cuda:7
10+
11+
cpus_per_task: 16
12+
eval_name: video_classification_frozen
13+
folder: /your_folder/evals/vitl/k400-area-attn
14+
mem_per_gpu: 220G
15+
nodes: 8
16+
num_workers: 8
17+
resume_checkpoint: true
18+
tag: k400-vitl16-16x8x3-16f-area-attn
19+
tasks_per_node: 8
20+
experiment:
21+
classifier:
22+
num_heads: 16
23+
num_probe_blocks: 4
24+
data:
25+
dataset_type: VideoDataset
26+
dataset_train: /your_data_path/k400_train_paths.csv
27+
dataset_val: /your_data_path/k400_val_paths.csv
28+
frame_step: 4
29+
frames_per_clip: 16
30+
num_classes: 400
31+
num_segments: 8
32+
num_views_per_segment: 3
33+
resolution: 256
34+
optimization:
35+
batch_size: 4
36+
multihead_kwargs:
37+
- final_lr: 0.0
38+
final_weight_decay: 0.01
39+
lr: 0.005
40+
start_lr: 0.005
41+
warmup: 0.0
42+
weight_decay: 0.01
43+
- final_lr: 0.0
44+
final_weight_decay: 0.01
45+
lr: 0.003
46+
start_lr: 0.003
47+
warmup: 0.0
48+
weight_decay: 0.01
49+
- final_lr: 0.0
50+
final_weight_decay: 0.01
51+
lr: 0.001
52+
start_lr: 0.001
53+
warmup: 0.0
54+
weight_decay: 0.01
55+
- final_lr: 0.0
56+
final_weight_decay: 0.01
57+
lr: 0.0003
58+
start_lr: 0.0003
59+
warmup: 0.0
60+
weight_decay: 0.01
61+
- final_lr: 0.0
62+
final_weight_decay: 0.01
63+
lr: 0.0001
64+
start_lr: 0.0001
65+
warmup: 0.0
66+
weight_decay: 0.01
67+
- final_lr: 0.0
68+
final_weight_decay: 0.1
69+
lr: 0.005
70+
start_lr: 0.005
71+
warmup: 0.0
72+
weight_decay: 0.1
73+
- final_lr: 0.0
74+
final_weight_decay: 0.1
75+
lr: 0.003
76+
start_lr: 0.003
77+
warmup: 0.0
78+
weight_decay: 0.1
79+
- final_lr: 0.0
80+
final_weight_decay: 0.1
81+
lr: 0.001
82+
start_lr: 0.001
83+
warmup: 0.0
84+
weight_decay: 0.1
85+
- final_lr: 0.0
86+
final_weight_decay: 0.1
87+
lr: 0.0003
88+
start_lr: 0.0003
89+
warmup: 0.0
90+
weight_decay: 0.1
91+
- final_lr: 0.0
92+
final_weight_decay: 0.1
93+
lr: 0.0001
94+
start_lr: 0.0001
95+
warmup: 0.0
96+
weight_decay: 0.1
97+
- final_lr: 0.0
98+
final_weight_decay: 0.4
99+
lr: 0.005
100+
start_lr: 0.005
101+
warmup: 0.0
102+
weight_decay: 0.4
103+
- final_lr: 0.0
104+
final_weight_decay: 0.4
105+
lr: 0.003
106+
start_lr: 0.003
107+
warmup: 0.0
108+
weight_decay: 0.4
109+
- final_lr: 0.0
110+
final_weight_decay: 0.4
111+
lr: 0.001
112+
start_lr: 0.001
113+
warmup: 0.0
114+
weight_decay: 0.4
115+
- final_lr: 0.0
116+
final_weight_decay: 0.4
117+
lr: 0.0003
118+
start_lr: 0.0003
119+
warmup: 0.0
120+
weight_decay: 0.4
121+
- final_lr: 0.0
122+
final_weight_decay: 0.4
123+
lr: 0.0001
124+
start_lr: 0.0001
125+
warmup: 0.0
126+
weight_decay: 0.4
127+
- final_lr: 0.0
128+
final_weight_decay: 0.8
129+
lr: 0.005
130+
start_lr: 0.005
131+
warmup: 0.0
132+
weight_decay: 0.8
133+
- final_lr: 0.0
134+
final_weight_decay: 0.8
135+
lr: 0.003
136+
start_lr: 0.003
137+
warmup: 0.0
138+
weight_decay: 0.8
139+
- final_lr: 0.0
140+
final_weight_decay: 0.8
141+
lr: 0.001
142+
start_lr: 0.001
143+
warmup: 0.0
144+
weight_decay: 0.8
145+
- final_lr: 0.0
146+
final_weight_decay: 0.8
147+
lr: 0.0003
148+
start_lr: 0.0003
149+
warmup: 0.0
150+
weight_decay: 0.8
151+
- final_lr: 0.0
152+
final_weight_decay: 0.8
153+
lr: 0.0001
154+
start_lr: 0.0001
155+
warmup: 0.0
156+
weight_decay: 0.8
157+
num_epochs: 20
158+
use_bfloat16: true
159+
use_pos_embed: false
160+
model_kwargs:
161+
checkpoint: /your_vjepa2_checkpoints/vitl-area-attn.pt
162+
module_name: evals.video_classification_frozen.modelcustom.vit_encoder_multiclip
163+
pretrain_kwargs:
164+
encoder:
165+
checkpoint_key: target_encoder
166+
img_temporal_dim_size: null
167+
model_name: vit_large
168+
patch_size: 16
169+
tubelet_size: 2
170+
uniform_power: true
171+
use_rope: true
172+
# -- ST-A² configuration (must match pretraining)
173+
use_area_attention: true
174+
area_attention_layers:
175+
- 0
176+
- 18
177+
area_spatial_splits: 2
178+
area_temporal_splits: 2
179+
area_residual_scale: 1.0
180+
wrapper_kwargs:
181+
max_frames: 128
182+
use_pos_embed: false

0 commit comments

Comments
 (0)