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
|
from __future__ import annotations
import argparse
from pathlib import Path
import torch
from torch.optim import AdamW
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer
from .common import (
cosine_schedule,
dtype_for_device,
parameter_count,
read_json,
seed_everything,
)
from .io import load_feature_pair, select_rows
from .models import PrefixAdapter
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
p.add_argument("--manifest", default="artifacts/manifest.json")
p.add_argument("--vision", default="artifacts/vision.pt")
p.add_argument("--text", default="artifacts/text.pt")
p.add_argument("--output", default="artifacts/prefix.pt")
p.add_argument("--device", default="cuda:1")
p.add_argument("--steps", type=int, default=3_000)
p.add_argument("--batch-size", type=int, default=32)
p.add_argument("--prefix-length", type=int, default=8)
p.add_argument("--hidden-dim", type=int, default=2048)
p.add_argument("--max-length", type=int, default=48)
p.add_argument("--lr", type=float, default=3e-4)
p.add_argument("--warmup", type=int, default=200)
p.add_argument("--seed", type=int, default=20260728)
return p.parse_args()
def main() -> None:
args = parse_args()
seed_everything(args.seed)
manifest = read_json(args.manifest)
_, text, _, tlookup = load_feature_pair(args.vision, args.text)
train_rows = manifest["text_only_train"]
semantic = select_rows(text["features"], tlookup, train_rows)
captions_by_row = {
int(row): caption for row, caption in zip(text["rows"], text["captions"])
}
captions = [captions_by_row[int(row)] for row in train_rows]
tokenizer = AutoTokenizer.from_pretrained(text["model"])
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
dtype = dtype_for_device(args.device)
lm = AutoModelForCausalLM.from_pretrained(
text["model"], torch_dtype=dtype
).to(args.device)
lm.eval()
for parameter in lm.parameters():
parameter.requires_grad_(False)
lm_dim = lm.get_input_embeddings().embedding_dim
adapter = PrefixAdapter(
semantic_dim=semantic.shape[-1],
lm_dim=lm_dim,
prefix_length=args.prefix_length,
hidden_dim=args.hidden_dim,
).to(args.device)
print(f"Prefix adapter parameters: {parameter_count(adapter):,}")
optimizer = AdamW(adapter.parameters(), lr=args.lr, weight_decay=1e-4)
generator = torch.Generator().manual_seed(args.seed)
history = []
progress = tqdm(range(args.steps), desc="text-only prefix")
for step in progress:
ids = torch.randint(
len(semantic), (args.batch_size,), generator=generator
)
batch_captions = [captions[int(i)] for i in ids]
tokens = tokenizer(
batch_captions,
padding=True,
truncation=True,
max_length=args.max_length,
return_tensors="pt",
)
input_ids = tokens["input_ids"].to(args.device)
attention = tokens["attention_mask"].to(args.device)
prefix = adapter(semantic[ids].to(args.device)).to(dtype)
token_embeddings = lm.get_input_embeddings()(input_ids)
inputs_embeds = torch.cat([prefix, token_embeddings], dim=1)
prefix_attention = torch.ones(
prefix.shape[:2], dtype=attention.dtype, device=args.device
)
full_attention = torch.cat([prefix_attention, attention], dim=1)
labels = input_ids.clone()
labels[attention == 0] = -100
prefix_labels = torch.full(
prefix.shape[:2], -100, dtype=labels.dtype, device=args.device
)
full_labels = torch.cat([prefix_labels, labels], dim=1)
result = lm(
inputs_embeds=inputs_embeds,
attention_mask=full_attention,
labels=full_labels,
use_cache=False,
return_dict=True,
)
loss = result.loss
optimizer.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
optimizer.step()
scale = cosine_schedule(step, args.steps, args.warmup)
for group in optimizer.param_groups:
group["lr"] = args.lr * scale
if step % 20 == 0:
progress.set_postfix(loss=f"{loss.item():.3f}")
if step % 100 == 0 or step == args.steps - 1:
history.append({"step": step, "loss": float(loss.item())})
state = {
"config": adapter.config(),
"state_dict": adapter.state_dict(),
"text_model": text["model"],
"text_layer": text["layer"],
"args": vars(args),
"history": history,
"training": "text only; no image features or image-text pairs",
}
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
torch.save(state, args.output)
print(f"Wrote {args.output}")
if __name__ == "__main__":
main()
|
