Upload 4 files

pretrain.py

@@ -34,18 +34,11 @@ import pickle as pkl
 from sophia import SophiaG
 
 
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 
-# # constants
-
-# NUM_BATCHES = int(1e5)
-# BATCH_SIZE = 4
 GRADIENT_ACCUMULATE_EVERY = 4
 LEARNING_RATE = 1e-4
 VALIDATE_EVERY  = 100
 GENERATE_EVERY  = 500
-# GENERATE_LENGTH = 2048
-# SEQ_LEN = 4096
 
 
 parser = argparse.ArgumentParser()
@@ -65,9 +58,6 @@ parser.add_argument("--ckpt_dir", type=str, default='./ckpts/', help='Directory
 parser.add_argument("--model_name", type=str, default='finetune', help='Finetuned model name.')
 
 args = parser.parse_args()
-# rank = int(os.environ["RANK"])
-# local_rank = args.local_rank
-# is_master = local_rank == 0
 
 SEED = args.seed
 EPOCHS = args.epoch
@@ -86,14 +76,6 @@ POS_EMBED_USING = args.pos_embed
 model_name = args.model_name
 ckpt_dir = args.ckpt_dir
 
-# dist.init_process_group(backend='nccl')
-# torch.cuda.set_device(local_rank)
-# device = torch.device("cuda", local_rank)
-# world_size = torch.distributed.get_world_size()
-
-# seed_all(SEED + torch.distributed.get_rank())
-
-
 
 # helpers
 
@@ -127,27 +109,7 @@ model = PerformerLM(
 model = AutoregressiveWrapper(model)
 model.cuda()
 
-
-
 # prepare sc data
-
-class SCDataset(Dataset):
-    def __init__(self, data, label):
-        super().__init__()
-        self.data = data
-        self.label = label
-
-    def __getitem__(self, index):
-        rand_start = random.randint(0, self.data.shape[0]-1)
-        full_seq = self.data[rand_start].toarray()[0]
-        full_seq[full_seq > (CLASS - 2)] = CLASS - 2
-        full_seq = torch.from_numpy(full_seq).long()
-        full_seq = torch.cat((full_seq, torch.tensor([0]))).to(device)
-        seq_label = self.label[rand_start]
-        return full_seq, seq_label
-
-    def __len__(self):
-        return self.data.shape[0]
     
 class SCDatasetPretrain(Dataset):
     def __init__(self, data, seq_len):
@@ -169,19 +131,8 @@ class SCDatasetPretrain(Dataset):
 
     def __len__(self):
         return self.data.shape[0]
-    
 
 data = sc.read_h5ad(args.data_path)
-#data = data[:1000, :]
-# label_dict, label = np.unique(np.array(data.obs['cell_type']), return_inverse=True)  # Convert strings categorical to integrate categorical, and label_dict[label] can be restored
-# #store the label dict and label for prediction
-# with open('label_dict', 'wb') as fp:
-#     pkl.dump(label_dict, fp)
-# with open('label', 'wb') as fp:
-#     pkl.dump(label, fp)
-# class_num = np.unique(label, return_counts=True)[1].tolist()
-# class_weight = torch.tensor([(1 - (x / sum(class_num))) ** 2 for x in class_num])
-# label = torch.from_numpy(label)
 data = data.X
 
 acc = []
@@ -190,18 +141,6 @@ f1w = []
 skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED)
 pred_list = pd.Series(['un'] * data.shape[0])
 
-# sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=SEED)
-# for index_train in sss.split(data):
-#     data_train = data[index_train]
-#     data_val = data[index_val]
-#     train_dataset = SCDatasetPretrain(data_train, SEQ_LEN)
-#     val_dataset = SCDatasetPretrain(data_val, SEQ_LEN)
-
-# train_sampler = DistributedSampler(train_dataset)
-# val_sampler = DistributedSampler(val_dataset)
-# train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=train_sampler)
-# val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, sampler=val_sampler)
-
 index_train = int(data.shape[0]*0.8)
 data_train = data[:index_train]
 data_val = data[index_train:]
@@ -210,15 +149,11 @@ val_dataset = SCDatasetPretrain(data_val, SEQ_LEN)
 
 train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE)
 val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE)
-# train_loader  = cycle(DataLoader(train_dataset, batch_size = BATCH_SIZE))
-# val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE))
 
 # optimizer
 
 optim = SophiaG(model.parameters(), lr=2e-4, 
                     betas=(0.965, 0.99), rho = 0.01, weight_decay=1e-1)
-# optim = torch.optim.SGD(model.parameters(), lr=1e-8, momentum=0.9)
-# optim = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
 scaler = GradScaler()
 
 # training
@@ -244,14 +179,6 @@ for i in tqdm(range(EPOCHS), mininterval=10., desc='training'):
     scaler.update()
     optim.zero_grad()
 
-    # if i % VALIDATE_EVERY == 0:
-    #     model.eval()
-    #     with torch.no_grad():
-    #         #loss = model(next(val_loader), return_loss = True)
-    #         for index, data_batch in enumerate(tqdm(val_loader)):
-    #             loss = model(data_batch, return_loss = True)
-    #             print(f'validation loss: {loss.item()}')
-
     if i % GENERATE_EVERY == 0 and i != 0:
         model.eval()
         inp = random.choice(val_dataset)[:-1]
@@ -266,5 +193,3 @@ for i in tqdm(range(EPOCHS), mininterval=10., desc='training'):
 print('save model')
 checkpoint = {'state_dict': model.state_dict(),'optimizer' :optim.state_dict()}
 torch.save(checkpoint, os.path.join(ckpt_dir, 'model_gene_attn.pth'))
-
-a=1

requirements.txt

@@ -0,0 +1,8 @@
+torch==1.8.1
+torchvision==0.9.1
+transformers==4.6.1
+scanpy==1.7.2
+scikit-learn==0.24.2
+scipy==1.5.4
+numpy==1.19.2
+pandas==1.1.5