TissueFormer/test_time_memory.py at main · uhlerlab/TissueFormer · GitHub

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
import torch
import torch.nn as nn
import torch.optim as optim
from torch_geometric.loader import DataLoader
import numpy as np
import pandas as pd
import random
import argparse
import pickle

from run_pretrain import run_pretrain
from run_finetune import run_update
from utils import dataset_create, dataset_create_split
from parse import parse_pretrain_method, parse_regression_method, parser_add_main_args
import os
import warnings
import time

os.environ["KMP_WARNINGS"] = "off"
warnings.filterwarnings('ignore')

def fix_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

import subprocess
def get_gpu_memory_map(device):
    result = subprocess.check_output(
        [
            'nvidia-smi', '--query-gpu=memory.used',
            '--format=csv,nounits,noheader'
        ], encoding='utf-8')
    gpu_memory = np.array([int(x) for x in result.strip().split('\n')])
    return gpu_memory[device]

### Parse args ###
parser = argparse.ArgumentParser(description='Pretraining Pipeline')
parser_add_main_args(parser)
args = parser.parse_args()

fix_seed(args.seed)

dir_path = '/data/wuqitian/hest_data_xenium_protein_preprocess'
meta_info = pd.read_csv("../data/meta_info_xenium.csv")
meta_info = meta_info[meta_info['tech'] == 'Xenium']
pretrain_samples = ['TENX126', 'TENX125', 'TENX124', 'TENX123', 'TENX122', 'TENX121', 'TENX119', 'TENX118',
                    'TENX110', 'TENX111', 'TENX139', 'TENX96', 'TENX99', 'TENX138', 'TENX95', 'TENX140', 'TENX98', 'TENX97']
sample = pretrain_samples

pretrain_dataset = dataset_create_split(dir_path, sample, args, valid_prop=0., test_prop=0.1)

print(len(pretrain_samples))

pretrain_dataloader = DataLoader(pretrain_dataset, batch_size=1, shuffle=True)

if args.cpu:
    device = torch.device("cpu")
else:
    device = torch.device("cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu")

# human protein-coding genes as shared gene idx
args.gene_total_num = 23258

# image emb dim of foundation encoders
if args.image_model == 'hoptimus':
    args.image_emb_dim = 1536
elif args.image_model == 'gigapath':
    args.image_emb_dim = 1536
elif args.image_model == 'uni':
    args.image_emb_dim = 1024

if args.gene_encoder_pretrained:
    gene_embeddings = pickle.load(open('../data/gene_embeddings_scgpt_human_protein.pkl', 'rb'))
else:
    gene_embeddings = torch.zeros((args.gene_total_num, args.gene_emb_dim), dtype=torch.float)
gene_embeddings = gene_embeddings.to(device)

model = parse_pretrain_method(args, gene_embeddings, device)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr_pretrain, weight_decay=args.wd_pretrain)

sizes, times, mems = [], [], []

model.train()
end_time = time.time()
for i, dataset in enumerate(pretrain_dataloader):
    data = dataset[0]
    x1, x2 = data.x.to(device), data.y.to(device)
    edge_index = data.edge_index.to(device)
    gene_idx = data.gene_idx.to(device)
    train_idx = data.split_idx['train_idx'].to(device)

    start_time = end_time
    loss = model.loss(x1, x2, gene_idx, train_idx, edge_index)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    end_time = time.time()
    memory_cost = get_gpu_memory_map(args.device)
    training_time = end_time - start_time

    print(f'{i}, sample size: {x1.shape[0]}, train time {training_time:.4f}, train memory {memory_cost:.1f}')
    sizes += [x1.shape[0]]
    times += [training_time]
    mems += [memory_cost]
    torch.cuda.empty_cache()

# model.eval()
# end_time = time.time()
# for i, dataset in enumerate(pretrain_dataloader):
#     data = dataset[0]
#     x1, x2 = data.x.to(device), data.y.to(device)
#     edge_index = data.edge_index.to(device)
#     gene_idx = data.gene_idx.to(device)
#     train_idx = data.split_idx['train_idx'].to(device)
#
#     start_time = end_time
#     outputs = model.encoder1(x1, edge_index)
#     end_time = time.time()
#     memory_cost = get_gpu_memory_map(args.device)
#     inference_time = end_time - start_time
#
#     print(f'{i}, sample size: {x1.shape[0]}, prediction time {inference_time:.4f}, prediction memory {memory_cost:.1f}')
#     sizes += [x1.shape[0]]
#     times += [inference_time]
#     mems += [memory_cost]
#     torch.cuda.empty_cache()

print(sizes)
print(times)
print(mems)