causal=False, # auto-regressive or not bucket_size=64, # average size of qk per bucket, 64 was recommended in paper n_hashes=4, # 4 is permissible per author, 8 is the best but slower ff_chunks= 200, # number of chunks for feedforward layer, make higher if there are memory issues weight_tie= False, # tie parameters of each layer for no memory per additional depth attn_chunks= 8, # process lsh attention in chunks, only way for memory to fit when scaling to 16k tokens num_mem_kv= 128, # persistent learned memory key values, from all-attention paper twin_attention= False, # both branches of the reversible network will be attention use_full_attn=False, # use full self attention, for comparison full_attn_thres= 128, # use full attention if context length is less than set value use_scale_norm= True, # use scale norm from 'Transformers without tears' paper axial_position_emb=True, axial_position_shape=(640, 64), axial_position_dims=(384, 384)) model.train() model.to(devices) print("starting test") inputs = torch.randint(low=0, high=tokenizer.vocab_size - 1, size=(10, tokenizer.max_len)).to(devices) output = model(inputs) print(output) print("test pass")
pred = model(inputs) pred.shape tokenizer.decode(torch.argmax(pred, dim=-1).squeeze(0)) loss_fn = nn.CrossEntropyLoss() # masked_lm_loss = loss_fn(pred.view(-1, tokenizer.vocab_size), labels.view(-1)) masked_lm_loss device = 'cuda:0' if torch.cuda.is_available() else 'cpu' total_loss = 0.0 model.train() model.to(device) inputs = inputs.to(device) labels = labels.to(device) loss = [] optimizer = AdamW(params=model.parameters()) for _ in tqdm(range(100000)): pred = model(inputs) mlm_loss = loss_fn(pred.view(-1, tokenizer.vocab_size), labels.view(-1)) total_loss += mlm_loss.item() loss.append(mlm_loss.item()) mlm_loss.backward() optimizer.step()
def test_encdec_v1(input_lang, target_lang, dim, bucket_size, depth, heads, n_hashes, vir_seq_len, ff_chunks, attn_chunks, mol_seq_len, cmd_args, train_dataset, test_dataset, output_folder, train_batch_size, epochs, validate_every, save_every, checkpoint_id, deepspeed_optimizer, use_full_attn, gradient_accumulation_steps, filter_thres): results = { 'generated_seq': [], 'generated_mol': [], 'target_mol': [], 'input_genome': [] } encoder = ReformerLM( num_tokens=input_lang.n_words, dim=dim, bucket_size=bucket_size, depth=depth, heads=heads, n_hashes=n_hashes, max_seq_len=vir_seq_len, ff_chunks=ff_chunks, attn_chunks=attn_chunks, weight_tie=True, weight_tie_embedding=True, axial_position_emb=True, axial_position_shape=compute_axial_position_shape(vir_seq_len), axial_position_dims=(dim // 2, dim // 2), return_embeddings=True, use_full_attn=use_full_attn).to(device) decoder = ReformerLM( num_tokens=target_lang.n_words, dim=dim, bucket_size=bucket_size, depth=depth, heads=heads, n_hashes=n_hashes, ff_chunks=ff_chunks, attn_chunks=attn_chunks, max_seq_len=mol_seq_len, axial_position_emb=True, axial_position_shape=compute_axial_position_shape(mol_seq_len), axial_position_dims=(dim // 2, dim // 2), weight_tie=True, weight_tie_embedding=True, causal=True, use_full_attn=use_full_attn).to(device) SAVE_DIR = os.sep.join([output_folder, 'saved_model']) if checkpoint_id: enc_ckp_max = checkpoint_id dec_ckp_max = checkpoint_id else: try: enc_ckp_max = np.max([ int(ckp) for ckp in os.listdir(os.sep.join([SAVE_DIR, 'encoder'])) ]) except Exception as e: print('Exception:', e) enc_ckp_max = 0 try: dec_ckp_max = np.max([ int(ckp) for ckp in os.listdir(os.sep.join([SAVE_DIR, 'decoder'])) ]) except: dec_ckp_max = 0 encoder = TrainingWrapper(encoder, ignore_index=PAD_IDX, pad_value=PAD_IDX).to(device) decoder = TrainingWrapper(decoder, ignore_index=PAD_IDX, pad_value=PAD_IDX).to(device) ''' encoder_params = filter(lambda p: p.requires_grad, encoder.parameters()) decoder_params = filter(lambda p: p.requires_grad, decoder.parameters()) if deepspeed_optimizer == False: print('No DeepSpeed optimizer found. Using RangerLars.') encoder_optimizer = RangerLars(encoder.parameters()) decoder_optimizer = RangerLars(decoder.parameters()) encoder_engine, encoder_optimizer, trainloader, _ = deepspeed.initialize( args=cmd_args, model=encoder, optimizer=encoder_optimizer, model_parameters=encoder_params, training_data=train_dataset, dist_init_required=True ) decoder_engine, decoder_optimizer, testloader, _ = deepspeed.initialize( args=cmd_args, model=decoder, optimizer=decoder_optimizer, model_parameters=decoder_params, training_data=test_dataset, dist_init_required=False ) else: print('Found optimizer in the DeepSpeed configurations. Using it.') encoder_engine, encoder_optimizer, trainloader, _ = deepspeed.initialize(args=cmd_args, model=encoder, model_parameters=encoder_params, training_data=train_dataset, dist_init_required=True) decoder_engine, decoder_optimizer, testloader, _ = deepspeed.initialize(args=cmd_args, model=decoder, model_parameters=decoder_params, training_data=test_dataset, dist_init_required=False) _, encoder_client_sd = encoder_engine.load_checkpoint(os.sep.join([SAVE_DIR,'encoder']), enc_ckp_max) _, decoder_client_sd = decoder_engine.load_checkpoint(os.sep.join([SAVE_DIR,'decoder']), dec_ckp_max) gpus_mini_batch = (train_batch_size// gradient_accumulation_steps) // torch.cuda.device_count() print('gpus_mini_batch:', gpus_mini_batch, 'with gradient_accumulation_steps:', gradient_accumulation_steps) for pair in tqdm(testloader): encoder_engine.eval() decoder_engine.eval() encoder.eval() decoder.eval() with torch.no_grad(): ts_src = pair[0] ts_trg = pair[1] input_genome = [[input_lang.index2word[gen_idx.item()] for gen_idx in smpl] for smpl in pair[0]] target_mol = [[target_lang.index2word[mol_idx.item()] for mol_idx in smpl] for smpl in pair[1]] ts_src = ts_src.to(encoder_engine.local_rank) #ts_src.to(device) # ts_trg = ts_trg.to(decoder_engine.local_rank) #ts_trg.to(device) # print('ts_src.shape', ts_src.shape) print('ts_src.shape', ts_trg.shape) enc_keys = encoder(ts_src) #encoder_engine(ts_src) yi = torch.tensor([[SOS_token] for _ in range(gpus_mini_batch)]).long().to(decoder_engine.local_rank) #to(device) # #sample = decoder_engine.generate(yi, mol_seq_len, filter_logits_fn=top_p, filter_thres=0.95, keys=enc_keys, eos_token = EOS_token) sample = decoder.generate(yi, mol_seq_len, filter_logits_fn=top_p, filter_thres=0.95, keys=enc_keys, eos_token = EOS_token) actual_mol = [] for mol_seq in sample.cpu().numpy(): for mol_idx in mol_seq: actual_mol.append(target_lang.index2word[mol_idx]) print('Generated Seq:', sample) print('Generated Mol:', actual_mol) print('Real Mol:', target_mol[:target_mol.index(target_lang.index2word[EOS_token])]) results['generated_seq'].append(sample) results['generated_mol'].append(actual_mol) results['target_mol'].append(target_mol) results['input_genome'].append(input_genome) print('Saving Test Results..') pickle.dump(results, open(os.sep.join([output_folder,'test_results.pkl']), 'wb')) ''' encoder_checkpoint = os.sep.join([ output_folder, 'saved_model', 'encoder', enc_ckp_max, 'mp_rank_00_model_states.pt' ]) decoder_checkpoint = os.sep.join([ output_folder, 'saved_model', 'decoder', dec_ckp_max, 'mp_rank_00_model_states.pt' ]) encoder.load_state_dict( torch.load(encoder_checkpoint, map_location=torch.device(device))['module']) decoder.load_state_dict( torch.load(decoder_checkpoint, map_location=torch.device(device))['module']) real_batch_size = train_batch_size // gradient_accumulation_steps test_loader = DataLoader(dataset=test_dataset, batch_size=real_batch_size, shuffle=True) if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs encoder = nn.DataParallel(encoder) decoder = nn.DataParallel(decoder) encoder.to(device) decoder.to(device) for pair in tqdm(test_loader): encoder.eval() decoder.eval() with torch.no_grad(): ts_src = torch.tensor(np.array([pair[0].numpy()])).to(device) ts_trg = torch.tensor(np.array([pair[1].numpy()])).to(device) input_genome = [ input_lang.index2word[gen_idx.item()] for gen_idx in pair[0] ] target_mol = [ target_lang.index2word[mol_idx.item()] for mol_idx in pair[1] ] enc_keys = encoder(ts_src) yi = torch.tensor([[SOS_token]]).long().to(device) sample = decoder.generate(yi, mol_seq_len, filter_logits_fn=top_p, filter_thres=filter_thres, keys=enc_keys, eos_token=EOS_token) actual_mol = [] for mol_seq in sample.cpu().numpy(): for mol_idx in mol_seq: actual_mol.append(target_lang.index2word[mol_idx]) print('Generated Seq:', sample) print('Generated Mol:', actual_mol) print( 'Real Mol:', target_mol[:target_mol.index(target_lang. index2word[EOS_token])]) results['generated_seq'].append(sample) results['generated_mol'].append(actual_mol) results['target_mol'].append(target_mol) results['input_genome'].append(input_genome) print('Saving Test Results..') pickle.dump(results, open(os.sep.join([output_folder, 'test_results.pkl']), 'wb')) '''