-
Notifications
You must be signed in to change notification settings - Fork 2
/
train_tangle.py
296 lines (238 loc) · 11 KB
/
train_tangle.py
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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
# --> General imports
import os
import numpy as np
from tqdm import tqdm
import json
# --> Torch imports
import torch
from torch.utils.data import DataLoader
import time
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.optim.lr_scheduler import LinearLR
# --> internal imports
from core.models.mmssl import MMSSL
from core.dataset.dataset import TangleDataset
from core.loss.tangle_loss import InfoNCE, apply_random_mask, init_intra_wsi_loss_function
from core.utils.learning import smooth_rank_measure, collate_tangle, set_seed
from core.utils.process_args import process_args
import pdb
# Set device
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def train_loop(args, loss_fn_interMod, loss_fn_rnaRecon, loss_fn_intraMod, ssl_model, epoch, dataloader, optimizer, scheduler_warmup, scheduler):
ssl_model.train()
ssl_model.to(DEVICE)
ep_loss, ep_recon_loss, ep_inter_loss, ep_intra_loss = 0., 0., 0., 0.
fb_time = 0.
all_embeds = []
for b_idx, (patch_emb, rna_seq, patch_emb_aug, avg_patch_emb) in enumerate(dataloader):
losses = []
s_fb = time.time()
# preprocessing for intra-modality loss
if args["intra_modality_wsi"]:
if args["intra_modality_mode_wsi"] == "contrast_token_views":
patch_emb = torch.cat((patch_emb, patch_emb_aug))
elif args["intra_modality_mode_wsi"] == "reconstruct_masked_emb" or args["intra_modality_mode_wsi"] == "reconstruct_masked_emb+contrast_avg_emb":
patch_emb_mask = apply_random_mask(patch_embeddings=patch_emb, percentage=args['mask_percentage'])
patch_emb = torch.cat((patch_emb, patch_emb_mask))
# set data on device
patch_emb = patch_emb.to(DEVICE)
rna_seq = rna_seq.to(DEVICE) if rna_seq is not None else rna_seq
if args["intra_modality_mode_wsi"] == "contrast_avg_emb" or args["intra_modality_mode_wsi"] == "reconstruct_avg_emb" or args["intra_modality_mode_wsi"] == "reconstruct_masked_emb+contrast_avg_emb":
avg_patch_emb = avg_patch_emb.cuda()
# forward pass and loss
if args["intra_modality_wsi"]:
wsi_emb, rna_emb, rna_reconstruction = ssl_model(patch_emb, None)
else:
wsi_emb, rna_emb, rna_reconstruction = ssl_model(patch_emb, rna_seq)
# intra modality loss wsi <-> wsi
if rna_emb is None and rna_reconstruction is None:
if args["intra_modality_mode_wsi"] == "contrast_token_views":
split_idx = int(patch_emb.shape[0]/2)
losses.append(loss_fn_intraMod(query=wsi_emb[:split_idx], positive_key=wsi_emb[split_idx:], symmetric=args["symmetric_cl"])) # 1. first set of token views 2. second set of token views (augmentation)
elif args["intra_modality_mode_wsi"] == "contrast_avg_emb":
losses.append(loss_fn_intraMod(query=wsi_emb, positive_key=avg_patch_emb, symmetric=args["symmetric_cl"]))
elif args["intra_modality_mode_wsi"] == "reconstruct_avg_emb":
losses.append(loss_fn_intraMod(wsi_emb, avg_patch_emb))
elif args["intra_modality_mode_wsi"] == "reconstruct_masked_emb":
split_idx = int(patch_emb.shape[0]/2)
losses.append(loss_fn_intraMod(wsi_emb[split_idx:], wsi_emb[:split_idx])) # 1. masked wsi_emb 2. umasked wsi_emb
elif args["intra_modality_mode_wsi"] == "reconstruct_masked_emb+contrast_avg_emb":
split_idx = int(patch_emb.shape[0]/2)
losses.append(loss_fn_intraMod(wsi_emb[split_idx:], wsi_emb[:split_idx])) # 1. masked wsi_emb 2. umasked wsi_emb
losses.append(loss_fn_intraMod(query=wsi_emb[:split_idx], positive_key=avg_patch_emb, symmetric=args["symmetric_cl"]))
else:
raise ValueError("Invalid intra_modality_mode_wsi.")
ep_intra_loss += losses[-1].item()
# inter modality loss wsi <-> rna
if rna_emb is not None:
losses.append(loss_fn_interMod(query=wsi_emb, positive_key=rna_emb, symmetric=args["symmetric_cl"]))
ep_inter_loss += losses[-1].item()
# intra modality loss rna <-> rna
if rna_reconstruction is not None:
losses.append(loss_fn_rnaRecon(rna_reconstruction, rna_seq))
ep_recon_loss += losses[-1].item()
loss = sum(losses)
optimizer.zero_grad()
loss.backward()
optimizer.step()
e_fb = time.time()
fb_time += e_fb - s_fb
if epoch <= args["warmup_epochs"]:
scheduler_warmup.step()
else:
scheduler.step()
if (b_idx % 3) == 0:
print(f"Loss for batch: {b_idx} = {loss}")
ep_loss += loss.item()
# get the train embeds to calculate rank
ssl_model.eval()
# do everything without grads
with torch.no_grad():
wsi_emb_to_store, _, _ = ssl_model(patch_emb)
all_embeds.extend(wsi_emb_to_store.detach().cpu().numpy())
ssl_model.train()
# track rank
all_embeds_tensor = torch.Tensor(np.array(all_embeds))
rank = smooth_rank_measure(all_embeds_tensor)
return ep_loss, rank
def val_loop(ssl_model, val_dataloader):
# set model to eval
ssl_model.eval()
ssl_model.to(DEVICE)
all_embeds = []
all_labels = []
# do everything without grads
with torch.no_grad():
for inputs, labels in tqdm(val_dataloader):
inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
wsi_embed, _, _ = ssl_model(inputs)
wsi_embed = wsi_embed.detach().cpu().numpy()
all_embeds.extend(wsi_embed)
all_labels.append(labels.item())
all_embeds = np.array(all_embeds)
all_labels = np.array(all_labels)
all_embeds_tensor = torch.Tensor(np.array(all_embeds))
rank = smooth_rank_measure(all_embeds_tensor)
results_dict = {"embeds": all_embeds, "labels": all_labels}
return results_dict, rank
def write_dict_to_config_file(config_dict, json_file_path):
"""
Write a dictionary to a configuration file.
Args:
config_dict (dict): The dictionary to be written to the config file.
config_file_path (str): The path to the configuration file.
Returns:
None
"""
with open(json_file_path, 'w') as jsonfile:
json.dump(config_dict, jsonfile, indent=4)
if __name__ == "__main__":
# setup args and seed
args = process_args()
args = vars(args)
set_seed(args["seed"])
# Set params for loss computation
RNA_RECONSTRUCTION = True if args["method"] == 'tanglerec' else False
INTRA_MODALITY = True if args["method"] == 'intra' else False
STOPPING_CRITERIA = 'train_rank' if args["method"] == 'tangle' or args["method"] == 'intra' else 'fixed'
N_TOKENS_RNA = 4908 if args["study"]=='nsclc' else 4999
args["rna_reconstruction"] = RNA_RECONSTRUCTION
args["intra_modality_wsi"] = INTRA_MODALITY
args["rna_token_dim"] = N_TOKENS_RNA
# paths
ROOT_SAVE_DIR = "./results/{}_checkpoints_and_embeddings".format(args["study"])
EXP_CODE = "{}_{}_lr{}_epochs{}_bs{}_tokensize{}_temperature{}_uni".format(
args["method"],
args["study"],
args["learning_rate"],
args["epochs"],
args["batch_size"],
args["n_tokens"],
args["temperature"]
)
RESULTS_SAVE_PATH = os.path.join(ROOT_SAVE_DIR, EXP_CODE)
os.makedirs(RESULTS_SAVE_PATH, exist_ok=True)
write_dict_to_config_file(args, os.path.join(RESULTS_SAVE_PATH, "config.json"))
print()
print(f"Running experiment {EXP_CODE}...")
print()
# Create a SummaryWriter
log_dir = os.path.join(ROOT_SAVE_DIR, 'logs', EXP_CODE)
os.makedirs(log_dir, exist_ok=True)
# make tangle dataset
print("* Setup dataset...")
dataset = TangleDataset(
feats_dir="./data/{}/uni_features/tcga_features/".format(args["study"]),
rna_dir='./data/{}/rna'.format(args["study"]),
sampling_strategy=args["sampling_strategy"],
n_tokens=args["n_tokens"]
)
# set up dataloader
print("* Setup dataloader...")
dataloader = DataLoader(
dataset,
batch_size=args["batch_size"],
shuffle=True,
collate_fn=collate_tangle
)
# set up model config, n_tokens_wsi, n_tokens_rna, patch_embedding_dim=768
print("* Setup model...")
ssl_model = MMSSL(config=args, n_tokens_rna=N_TOKENS_RNA).to(DEVICE)
if len(args["gpu_devices"]) > 1:
print(f"* Using {torch.cuda.device_count()} GPUs.")
ssl_model = nn.DataParallel(ssl_model, device_ids=args["gpu_devices"])
ssl_model.to("cuda:0")
# set up optimizers
print("* Setup optimizer...")
optimizer = optim.AdamW(ssl_model.parameters(), lr=args["learning_rate"])
# set up schedulers
print("* Setup schedulers...")
T_max = (args["epochs"] - args["warmup_epochs"]) * len(dataloader) if args["warmup"] else args["epochs"] * len(dataloader)
scheduler = CosineAnnealingLR(
optimizer,
T_max=T_max,
eta_min=args["end_learning_rate"]
)
if args["warmup"]:
scheduler_warmup = LinearLR(
optimizer,
start_factor=0.00001,
total_iters=args["warmup_epochs"] * len(dataloader)
)
else:
scheduler_warmup = None
# set up losses
print("* Setup losses...")
loss_fn_interMod = InfoNCE(temperature=args["temperature"])
loss_fn_rnaRecon = nn.MSELoss()
loss_fn_intraMod = init_intra_wsi_loss_function(args)
# main training loop
best_rank = 0.
for epoch in range(args["epochs"]):
print()
print(f"Training for epoch {epoch}...")
print()
# train
start = time.time()
ep_loss, train_rank = train_loop(args, loss_fn_interMod, loss_fn_rnaRecon, loss_fn_intraMod, ssl_model, epoch, dataloader, optimizer, scheduler_warmup, scheduler)
end = time.time()
print()
print(f"Done with epoch {epoch}")
print(f"Total loss = {ep_loss}")
print(f"Train rank = {train_rank}")
print("Total time = {:.3f} seconds".format(end-start))
# Stop training based on rank of the training samples. Ok for TANGLE and Intra.
if STOPPING_CRITERIA == 'train_rank':
if train_rank > best_rank:
print('Better rank: {} --> {}. Saving model'.format(best_rank, train_rank))
best_rank = train_rank
torch.save(ssl_model.state_dict(), os.path.join(RESULTS_SAVE_PATH, "model.pt"))
# Otherwise, stop after fixed number of training epochs. Ok for TANGLE-Rec.
else:
torch.save(ssl_model.state_dict(), os.path.join(RESULTS_SAVE_PATH, "model.pt"))
print()
print()
print("Done")
print()