import os
import time
import copy
from pathlib import Path
from math import ceil
from contextlib import contextmanager, nullcontext
from functools import partial, wraps
from collections.abc import Iterable
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import random_split, DataLoader
from torch.optim import Adam
from torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR
from torch.cuda.amp import autocast, GradScaler
import pytorch_warmup as warmup
import shutil
import esm
from einops import rearrange
from packaging import version
__version__ = '1.9.3'
device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu"
)
import matplotlib.pyplot as plt
def cycle(dl):
while True:
for data in dl:
yield data
from packaging import version
import numpy as np
from ema_pytorch import EMA
from accelerate import Accelerator, DistributedType, DistributedDataParallelKwargs
from fsspec.core import url_to_fs
from fsspec.implementations.local import LocalFileSystem
# # --
# from PD_SpLMxDiff.ModelPack import resize_image_to,ProteinDesigner_B
# from PD_SpLMxDiff.UtilityPack import get_Model_A_error, convert_into_tokens
# from PD_SpLMxDiff.UtilityPack import decode_one_ems_token_rec,decode_many_ems_token_rec
# from PD_SpLMxDiff.UtilityPack import decode_one_ems_token_rec_for_folding,decode_many_ems_token_rec_for_folding
# from PD_SpLMxDiff.UtilityPack import decode_one_ems_token_rec_for_folding_with_mask,decode_many_ems_token_rec_for_folding_with_mask,read_mask_from_input
# from PD_SpLMxDiff.UtilityPack import get_DSSP_result, string_diff
# from PD_SpLMxDiff.DataSetPack import pad_a_np_arr
# ++
from PD_pLMProbXDiff.ModelPack import (
resize_image_to, ProteinDesigner_B,
)
from PD_pLMProbXDiff.UtilityPack import (
get_Model_A_error, convert_into_tokens,convert_into_tokens_using_prob,
decode_one_ems_token_rec, decode_many_ems_token_rec,
decode_one_ems_token_rec_for_folding,
decode_many_ems_token_rec_for_folding,
decode_one_ems_token_rec_for_folding_with_mask,
decode_many_ems_token_rec_for_folding_with_mask,
read_mask_from_input,
get_DSSP_result,
string_diff,
load_in_pLM,
)
from PD_pLMProbXDiff.DataSetPack import (
pad_a_np_arr
)
# loss function
criterion_MSE_sum = nn.MSELoss(reduction='sum')
criterion_MAE_sum = nn.L1Loss(reduction='sum')
# helper functions
def exists(val):
return val is not None
def default(val, d):
if exists(val):
return val
return d() if callable(d) else d
def cast_tuple(val, length = 1):
if isinstance(val, list):
val = tuple(val)
return val if isinstance(val, tuple) else ((val,) * length)
def find_first(fn, arr):
for ind, el in enumerate(arr):
if fn(el):
return ind
return -1
def pick_and_pop(keys, d):
values = list(map(lambda key: d.pop(key), keys))
return dict(zip(keys, values))
def group_dict_by_key(cond, d):
return_val = [dict(),dict()]
for key in d.keys():
match = bool(cond(key))
ind = int(not match)
return_val[ind][key] = d[key]
return (*return_val,)
def string_begins_with(prefix, str):
return str.startswith(prefix)
def group_by_key_prefix(prefix, d):
return group_dict_by_key(partial(string_begins_with, prefix), d)
def groupby_prefix_and_trim(prefix, d):
kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, prefix), d)
kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items())))
return kwargs_without_prefix, kwargs
def num_to_groups(num, divisor):
groups = num // divisor
remainder = num % divisor
arr = [divisor] * groups
if remainder > 0:
arr.append(remainder)
return arr
# url to fs, bucket, path - for checkpointing to cloud
def url_to_bucket(url):
if '://' not in url:
return url
_, suffix = url.split('://')
if prefix in {'gs', 's3'}:
return suffix.split('/')[0]
else:
raise ValueError(f'storage type prefix "{prefix}" is not supported yet')
# decorators
def eval_decorator(fn):
def inner(model, *args, **kwargs):
was_training = model.training
model.eval()
out = fn(model, *args, **kwargs)
model.train(was_training)
return out
return inner
def cast_torch_tensor(fn, cast_fp16 = False):
@wraps(fn)
def inner(model, *args, **kwargs):
device = kwargs.pop('_device', model.device)
cast_device = kwargs.pop('_cast_device', True)
should_cast_fp16 = cast_fp16 and model.cast_half_at_training
kwargs_keys = kwargs.keys()
all_args = (*args, *kwargs.values())
split_kwargs_index = len(all_args) - len(kwargs_keys)
all_args = tuple(map(lambda t: torch.from_numpy(t) if exists(t) and isinstance(t, np.ndarray) else t, all_args))
if cast_device:
all_args = tuple(map(lambda t: t.to(device) if exists(t) and isinstance(t, torch.Tensor) else t, all_args))
if should_cast_fp16:
all_args = tuple(map(lambda t: t.half() if exists(t) and isinstance(t, torch.Tensor) and t.dtype != torch.bool else t, all_args))
args, kwargs_values = all_args[:split_kwargs_index], all_args[split_kwargs_index:]
kwargs = dict(tuple(zip(kwargs_keys, kwargs_values)))
out = fn(model, *args, **kwargs)
return out
return inner
# gradient accumulation functions
def split_iterable(it, split_size):
accum = []
for ind in range(ceil(len(it) / split_size)):
start_index = ind * split_size
accum.append(it[start_index: (start_index + split_size)])
return accum
def split(t, split_size = None):
if not exists(split_size):
return t
if isinstance(t, torch.Tensor):
return t.split(split_size, dim = 0)
if isinstance(t, Iterable):
return split_iterable(t, split_size)
return TypeError
def find_first(cond, arr):
for el in arr:
if cond(el):
return el
return None
def split_args_and_kwargs(*args, split_size = None, **kwargs):
all_args = (*args, *kwargs.values())
len_all_args = len(all_args)
first_tensor = find_first(lambda t: isinstance(t, torch.Tensor), all_args)
assert exists(first_tensor)
batch_size = len(first_tensor)
split_size = default(split_size, batch_size)
num_chunks = ceil(batch_size / split_size)
dict_len = len(kwargs)
dict_keys = kwargs.keys()
split_kwargs_index = len_all_args - dict_len
split_all_args = [split(arg, split_size = split_size) if exists(arg) and isinstance(arg, (torch.Tensor, Iterable)) else ((arg,) * num_chunks) for arg in all_args]
chunk_sizes = tuple(map(len, split_all_args[0]))
for (chunk_size, *chunked_all_args) in tuple(zip(chunk_sizes, *split_all_args)):
chunked_args, chunked_kwargs_values = chunked_all_args[:split_kwargs_index], chunked_all_args[split_kwargs_index:]
chunked_kwargs = dict(tuple(zip(dict_keys, chunked_kwargs_values)))
chunk_size_frac = chunk_size / batch_size
yield chunk_size_frac, (chunked_args, chunked_kwargs)
# imagen trainer
def imagen_sample_in_chunks(fn):
@wraps(fn)
def inner(self, *args, max_batch_size = None, **kwargs):
if not exists(max_batch_size):
return fn(self, *args, **kwargs)
if self.imagen.unconditional:
batch_size = kwargs.get('batch_size')
batch_sizes = num_to_groups(batch_size, max_batch_size)
outputs = [fn(self, *args, **{**kwargs, 'batch_size': sub_batch_size}) for sub_batch_size in batch_sizes]
else:
outputs = [fn(self, *chunked_args, **chunked_kwargs) for _, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs)]
if isinstance(outputs[0], torch.Tensor):
return torch.cat(outputs, dim = 0)
return list(map(lambda t: torch.cat(t, dim = 0), list(zip(*outputs))))
return inner
def restore_parts(state_dict_target, state_dict_from):
for name, param in state_dict_from.items():
if name not in state_dict_target:
continue
if param.size() == state_dict_target[name].size():
state_dict_target[name].copy_(param)
else:
print(f"layer {name}({param.size()} different than target: {state_dict_target[name].size()}")
return state_dict_target
class ImagenTrainer(nn.Module):
locked = False
def __init__(
self,
#imagen = None,
model = None,
imagen_checkpoint_path = None,
use_ema = True,
lr = 1e-4,
eps = 1e-8,
beta1 = 0.9,
beta2 = 0.99,
max_grad_norm = None,
group_wd_params = True,
warmup_steps = None,
cosine_decay_max_steps = None,
only_train_unet_number = None,
fp16 = False,
precision = None,
split_batches = True,
dl_tuple_output_keywords_names = ('images', 'text_embeds', 'text_masks', 'cond_images'),
verbose = True,
split_valid_fraction = 0.025,
split_valid_from_train = False,
split_random_seed = 42,
checkpoint_path = None,
checkpoint_every = None,
checkpoint_fs = None,
fs_kwargs: dict = None,
max_checkpoints_keep = 20,
# +++++++++++++++++++++
CKeys=None,
#
**kwargs
):
super().__init__()
assert not ImagenTrainer.locked, 'ImagenTrainer can only be initialized once per process - for the sake of distributed training, you will now have to create a separate script to train each unet (or a script that accepts unet number as an argument)'
assert exists(model.imagen) ^ exists(imagen_checkpoint_path), 'either imagen instance is passed into the trainer, or a checkpoint path that contains the imagen config'
# determine filesystem, using fsspec, for saving to local filesystem or cloud
self.fs = checkpoint_fs
if not exists(self.fs):
fs_kwargs = default(fs_kwargs, {})
self.fs, _ = url_to_fs(default(checkpoint_path, './'), **fs_kwargs)
# # -----------------------------------
# # from MJB
# assert isinstance(model.imagen, (ProteinDesigner_B))
# modified by BN
# ++: try this trainer for all models
# assert isinstance(model, (ProteinDesigner_B))
# +++++++++++++++++++++++++
self.CKeys = CKeys
ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)
self.imagen = model.imagen
self.model=model
self.is_elucidated = self.model.is_elucidated
# create accelerator instance
accelerate_kwargs, kwargs = groupby_prefix_and_trim('accelerate_', kwargs)
assert not (fp16 and exists(precision)), 'either set fp16 = True or forward the precision ("fp16", "bf16") to Accelerator'
accelerator_mixed_precision = default(precision, 'fp16' if fp16 else 'no')
self.accelerator = Accelerator(**{
'split_batches': split_batches,
'mixed_precision': accelerator_mixed_precision,
'kwargs_handlers': [DistributedDataParallelKwargs(find_unused_parameters = True)]
, **accelerate_kwargs})
ImagenTrainer.locked = self.is_distributed
# cast data to fp16 at training time if needed
self.cast_half_at_training = accelerator_mixed_precision == 'fp16'
# grad scaler must be managed outside of accelerator
grad_scaler_enabled = fp16
self.num_unets = len(self.imagen.unets)
self.use_ema = use_ema and self.is_main
self.ema_unets = nn.ModuleList([])
# keep track of what unet is being trained on
# only going to allow 1 unet training at a time
self.ema_unet_being_trained_index = -1 # keeps track of which ema unet is being trained on
# data related functions
self.train_dl_iter = None
self.train_dl = None
self.valid_dl_iter = None
self.valid_dl = None
self.dl_tuple_output_keywords_names = dl_tuple_output_keywords_names
# auto splitting validation from training, if dataset is passed in
self.split_valid_from_train = split_valid_from_train
assert 0 <= split_valid_fraction <= 1, 'split valid fraction must be between 0 and 1'
self.split_valid_fraction = split_valid_fraction
self.split_random_seed = split_random_seed
# be able to finely customize learning rate, weight decay
# per unet
lr, eps, warmup_steps, cosine_decay_max_steps = map(partial(cast_tuple, length = self.num_unets), (lr, eps, warmup_steps, cosine_decay_max_steps))
for ind, (unet, unet_lr, unet_eps, unet_warmup_steps, unet_cosine_decay_max_steps) in enumerate(zip(self.imagen.unets, lr, eps, warmup_steps, cosine_decay_max_steps)):
optimizer = Adam(
unet.parameters(),
lr = unet_lr,
eps = unet_eps,
betas = (beta1, beta2),
**kwargs
)
if self.use_ema:
self.ema_unets.append(EMA(unet, **ema_kwargs))
scaler = GradScaler(enabled = grad_scaler_enabled)
scheduler = warmup_scheduler = None
if exists(unet_cosine_decay_max_steps):
scheduler = CosineAnnealingLR(optimizer, T_max = unet_cosine_decay_max_steps)
if exists(unet_warmup_steps):
warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period = unet_warmup_steps)
if not exists(scheduler):
scheduler = LambdaLR(optimizer, lr_lambda = lambda step: 1.0)
# set on object
setattr(self, f'optim{ind}', optimizer) # cannot use pytorch ModuleList for some reason with optimizers
setattr(self, f'scaler{ind}', scaler)
setattr(self, f'scheduler{ind}', scheduler)
setattr(self, f'warmup{ind}', warmup_scheduler)
# gradient clipping if needed
self.max_grad_norm = max_grad_norm
# step tracker and misc
self.register_buffer('steps', torch.tensor([0] * self.num_unets))
self.verbose = verbose
# automatic set devices based on what accelerator decided
self.imagen.to(self.device)
self.to(self.device)
# checkpointing
assert not (exists(checkpoint_path) ^ exists(checkpoint_every))
self.checkpoint_path = checkpoint_path
self.checkpoint_every = checkpoint_every
self.max_checkpoints_keep = max_checkpoints_keep
self.can_checkpoint = self.is_local_main if isinstance(checkpoint_fs, LocalFileSystem) else self.is_main
if exists(checkpoint_path) and self.can_checkpoint:
bucket = url_to_bucket(checkpoint_path)
if not self.fs.exists(bucket):
self.fs.mkdir(bucket)
self.load_from_checkpoint_folder()
# only allowing training for unet
self.only_train_unet_number = only_train_unet_number
self.validate_and_set_unet_being_trained(only_train_unet_number)
# computed values
@property
def device(self):
return self.accelerator.device
@property
def is_distributed(self):
return not (self.accelerator.distributed_type == DistributedType.NO and self.accelerator.num_processes == 1)
@property
def is_main(self):
return self.accelerator.is_main_process
@property
def is_local_main(self):
return self.accelerator.is_local_main_process
@property
def unwrapped_unet(self):
return self.accelerator.unwrap_model(self.unet_being_trained)
# optimizer helper functions
def get_lr(self, unet_number):
self.validate_unet_number(unet_number)
unet_index = unet_number - 1
optim = getattr(self, f'optim{unet_index}')
return optim.param_groups[0]['lr']
# function for allowing only one unet from being trained at a time
def validate_and_set_unet_being_trained(self, unet_number = None):
if exists(unet_number):
self.validate_unet_number(unet_number)
assert not exists(self.only_train_unet_number) or self.only_train_unet_number == unet_number, 'you cannot only train on one unet at a time. you will need to save the trainer into a checkpoint, and resume training on a new unet'
self.only_train_unet_number = unet_number
self.imagen.only_train_unet_number = unet_number
if not exists(unet_number):
return
self.wrap_unet(unet_number)
def wrap_unet(self, unet_number):
if hasattr(self, 'one_unet_wrapped'):
return
unet = self.imagen.get_unet(unet_number)
self.unet_being_trained = self.accelerator.prepare(unet)
unet_index = unet_number - 1
optimizer = getattr(self, f'optim{unet_index}')
scheduler = getattr(self, f'scheduler{unet_index}')
optimizer = self.accelerator.prepare(optimizer)
if exists(scheduler):
scheduler = self.accelerator.prepare(scheduler)
setattr(self, f'optim{unet_index}', optimizer)
setattr(self, f'scheduler{unet_index}', scheduler)
self.one_unet_wrapped = True
# hacking accelerator due to not having separate gradscaler per optimizer
def set_accelerator_scaler(self, unet_number):
unet_number = self.validate_unet_number(unet_number)
scaler = getattr(self, f'scaler{unet_number - 1}')
self.accelerator.scaler = scaler
for optimizer in self.accelerator._optimizers:
optimizer.scaler = scaler
# helper print
def print(self, msg):
if not self.is_main:
return
if not self.verbose:
return
return self.accelerator.print(msg)
# validating the unet number
def validate_unet_number(self, unet_number = None):
if self.num_unets == 1:
unet_number = default(unet_number, 1)
assert 0 < unet_number <= self.num_unets, f'unet number should be in between 1 and {self.num_unets}'
return unet_number
# number of training steps taken
def num_steps_taken(self, unet_number = None):
if self.num_unets == 1:
unet_number = default(unet_number, 1)
return self.steps[unet_number - 1].item()
def print_untrained_unets(self):
print_final_error = False
for ind, (steps, unet) in enumerate(zip(self.steps.tolist(), self.imagen.unets)):
if steps > 0 or isinstance(unet, NullUnet):
continue
self.print(f'unet {ind + 1} has not been trained')
print_final_error = True
if print_final_error:
self.print('when sampling, you can pass stop_at_unet_number to stop early in the cascade, so it does not try to generate with untrained unets')
# data related functions
def add_train_dataloader(self, dl = None):
if not exists(dl):
return
assert not exists(self.train_dl), 'training dataloader was already added'
self.train_dl = self.accelerator.prepare(dl)
def add_valid_dataloader(self, dl):
if not exists(dl):
return
assert not exists(self.valid_dl), 'validation dataloader was already added'
self.valid_dl = self.accelerator.prepare(dl)
def add_train_dataset(self, ds = None, *, batch_size, **dl_kwargs):
if not exists(ds):
return
assert not exists(self.train_dl), 'training dataloader was already added'
valid_ds = None
if self.split_valid_from_train:
train_size = int((1 - self.split_valid_fraction) * len(ds))
valid_size = len(ds) - train_size
ds, valid_ds = random_split(ds, [train_size, valid_size], generator = torch.Generator().manual_seed(self.split_random_seed))
self.print(f'training with dataset of {len(ds)} samples and validating with randomly splitted {len(valid_ds)} samples')
dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
self.train_dl = self.accelerator.prepare(dl)
if not self.split_valid_from_train:
return
self.add_valid_dataset(valid_ds, batch_size = batch_size, **dl_kwargs)
def add_valid_dataset(self, ds, *, batch_size, **dl_kwargs):
if not exists(ds):
return
assert not exists(self.valid_dl), 'validation dataloader was already added'
dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
self.valid_dl = self.accelerator.prepare(dl)
def create_train_iter(self):
assert exists(self.train_dl), 'training dataloader has not been registered with the trainer yet'
if exists(self.train_dl_iter):
return
self.train_dl_iter = cycle(self.train_dl)
def create_valid_iter(self):
assert exists(self.valid_dl), 'validation dataloader has not been registered with the trainer yet'
if exists(self.valid_dl_iter):
return
self.valid_dl_iter = cycle(self.valid_dl)
def train_step(self, unet_number = None, **kwargs):
self.create_train_iter()
loss = self.step_with_dl_iter(self.train_dl_iter, unet_number = unet_number, **kwargs)
self.update(unet_number = unet_number)
return loss
@torch.no_grad()
@eval_decorator
def valid_step(self, **kwargs):
self.create_valid_iter()
context = self.use_ema_unets if kwargs.pop('use_ema_unets', False) else nullcontext
with context():
loss = self.step_with_dl_iter(self.valid_dl_iter, **kwargs)
return loss
def step_with_dl_iter(self, dl_iter, **kwargs):
dl_tuple_output = cast_tuple(next(dl_iter))
model_input = dict(list(zip(self.dl_tuple_output_keywords_names, dl_tuple_output)))
loss = self.forward(**{**kwargs, **model_input})
return loss
# checkpointing functions
@property
def all_checkpoints_sorted(self):
glob_pattern = os.path.join(self.checkpoint_path, '*.pt')
checkpoints = self.fs.glob(glob_pattern)
sorted_checkpoints = sorted(checkpoints, key = lambda x: int(str(x).split('.')[-2]), reverse = True)
return sorted_checkpoints
def load_from_checkpoint_folder(self, last_total_steps = -1):
if last_total_steps != -1:
filepath = os.path.join(self.checkpoint_path, f'checkpoint.{last_total_steps}.pt')
self.load(filepath)
return
sorted_checkpoints = self.all_checkpoints_sorted
if len(sorted_checkpoints) == 0:
self.print(f'no checkpoints found to load from at {self.checkpoint_path}')
return
last_checkpoint = sorted_checkpoints[0]
self.load(last_checkpoint)
def save_to_checkpoint_folder(self):
self.accelerator.wait_for_everyone()
if not self.can_checkpoint:
return
total_steps = int(self.steps.sum().item())
filepath = os.path.join(self.checkpoint_path, f'checkpoint.{total_steps}.pt')
self.save(filepath)
if self.max_checkpoints_keep <= 0:
return
sorted_checkpoints = self.all_checkpoints_sorted
checkpoints_to_discard = sorted_checkpoints[self.max_checkpoints_keep:]
for checkpoint in checkpoints_to_discard:
self.fs.rm(checkpoint)
# saving and loading functions
def save(
self,
path,
overwrite = True,
without_optim_and_sched = False,
**kwargs
):
self.accelerator.wait_for_everyone()
if not self.can_checkpoint:
return
fs = self.fs
assert not (fs.exists(path) and not overwrite)
self.reset_ema_unets_all_one_device()
save_obj = dict(
model = self.imagen.state_dict(),
version = __version__,
steps = self.steps.cpu(),
**kwargs
)
save_optim_and_sched_iter = range(0, self.num_unets) if not without_optim_and_sched else tuple()
for ind in save_optim_and_sched_iter:
scaler_key = f'scaler{ind}'
optimizer_key = f'optim{ind}'
scheduler_key = f'scheduler{ind}'
warmup_scheduler_key = f'warmup{ind}'
scaler = getattr(self, scaler_key)
optimizer = getattr(self, optimizer_key)
scheduler = getattr(self, scheduler_key)
warmup_scheduler = getattr(self, warmup_scheduler_key)
if exists(scheduler):
save_obj = {**save_obj, scheduler_key: scheduler.state_dict()}
if exists(warmup_scheduler):
save_obj = {**save_obj, warmup_scheduler_key: warmup_scheduler.state_dict()}
save_obj = {**save_obj, scaler_key: scaler.state_dict(), optimizer_key: optimizer.state_dict()}
if self.use_ema:
save_obj = {**save_obj, 'ema': self.ema_unets.state_dict()}
# determine if imagen config is available
if hasattr(self.imagen, '_config'):
self.print(f'this checkpoint is commandable from the CLI - "imagen --model {str(path)} \"<prompt>\""')
save_obj = {
**save_obj,
'imagen_type': 'elucidated' if self.is_elucidated else 'original',
'imagen_params': self.imagen._config
}
#save to path
with fs.open(path, 'wb') as f:
torch.save(save_obj, f)
self.print(f'checkpoint saved to {path}')
def load(self, path, only_model = False, strict = True, noop_if_not_exist = False):
fs = self.fs
if noop_if_not_exist and not fs.exists(path):
self.print(f'trainer checkpoint not found at {str(path)}')
return
assert fs.exists(path), f'{path} does not exist'
self.reset_ema_unets_all_one_device()
# to avoid extra GPU memory usage in main process when using Accelerate
with fs.open(path) as f:
loaded_obj = torch.load(f, map_location='cpu')
if version.parse(__version__) != version.parse(loaded_obj['version']):
self.print(f'loading saved imagen at version {loaded_obj["version"]}, but current package version is {__version__}')
try:
self.imagen.load_state_dict(loaded_obj['model'], strict = strict)
except RuntimeError:
print("Failed loading state dict. Trying partial load")
self.imagen.load_state_dict(restore_parts(self.imagen.state_dict(),
loaded_obj['model']))
if only_model:
return loaded_obj
self.steps.copy_(loaded_obj['steps'])
for ind in range(0, self.num_unets):
scaler_key = f'scaler{ind}'
optimizer_key = f'optim{ind}'
scheduler_key = f'scheduler{ind}'
warmup_scheduler_key = f'warmup{ind}'
scaler = getattr(self, scaler_key)
optimizer = getattr(self, optimizer_key)
scheduler = getattr(self, scheduler_key)
warmup_scheduler = getattr(self, warmup_scheduler_key)
if exists(scheduler) and scheduler_key in loaded_obj:
scheduler.load_state_dict(loaded_obj[scheduler_key])
if exists(warmup_scheduler) and warmup_scheduler_key in loaded_obj:
warmup_scheduler.load_state_dict(loaded_obj[warmup_scheduler_key])
if exists(optimizer):
try:
optimizer.load_state_dict(loaded_obj[optimizer_key])
scaler.load_state_dict(loaded_obj[scaler_key])
except:
self.print('could not load optimizer and scaler, possibly because you have turned on mixed precision training since the last run. resuming with new optimizer and scalers')
if self.use_ema:
assert 'ema' in loaded_obj
try:
self.ema_unets.load_state_dict(loaded_obj['ema'], strict = strict)
except RuntimeError:
print("Failed loading state dict. Trying partial load")
self.ema_unets.load_state_dict(restore_parts(self.ema_unets.state_dict(),
loaded_obj['ema']))
self.print(f'checkpoint loaded from {path}')
return loaded_obj
# managing ema unets and their devices
@property
def unets(self):
return nn.ModuleList([ema.ema_model for ema in self.ema_unets])
def get_ema_unet(self, unet_number = None):
if not self.use_ema:
return
unet_number = self.validate_unet_number(unet_number)
index = unet_number - 1
if isinstance(self.unets, nn.ModuleList):
unets_list = [unet for unet in self.ema_unets]
delattr(self, 'ema_unets')
self.ema_unets = unets_list
if index != self.ema_unet_being_trained_index:
for unet_index, unet in enumerate(self.ema_unets):
unet.to(self.device if unet_index == index else 'cpu')
self.ema_unet_being_trained_index = index
return self.ema_unets[index]
def reset_ema_unets_all_one_device(self, device = None):
if not self.use_ema:
return
device = default(device, self.device)
self.ema_unets = nn.ModuleList([*self.ema_unets])
self.ema_unets.to(device)
self.ema_unet_being_trained_index = -1
@torch.no_grad()
@contextmanager
def use_ema_unets(self):
if not self.use_ema:
output = yield
return output
self.reset_ema_unets_all_one_device()
self.imagen.reset_unets_all_one_device()
self.unets.eval()
trainable_unets = self.imagen.unets
self.imagen.unets = self.unets # swap in exponential moving averaged unets for sampling
output = yield
self.imagen.unets = trainable_unets # restore original training unets
# cast the ema_model unets back to original device
for ema in self.ema_unets:
ema.restore_ema_model_device()
return output
def print_unet_devices(self):
self.print('unet devices:')
for i, unet in enumerate(self.imagen.unets):
device = next(unet.parameters()).device
self.print(f'\tunet {i}: {device}')
if not self.use_ema:
return
self.print('\nema unet devices:')
for i, ema_unet in enumerate(self.ema_unets):
device = next(ema_unet.parameters()).device
self.print(f'\tema unet {i}: {device}')
# overriding state dict functions
def state_dict(self, *args, **kwargs):
self.reset_ema_unets_all_one_device()
return super().state_dict(*args, **kwargs)
def load_state_dict(self, *args, **kwargs):
self.reset_ema_unets_all_one_device()
return super().load_state_dict(*args, **kwargs)
# encoding text functions
def encode_text(self, text, **kwargs):
return self.imagen.encode_text(text, **kwargs)
# forwarding functions and gradient step updates
def update(self, unet_number = None):
unet_number = self.validate_unet_number(unet_number)
self.validate_and_set_unet_being_trained(unet_number)
self.set_accelerator_scaler(unet_number)
index = unet_number - 1
unet = self.unet_being_trained
optimizer = getattr(self, f'optim{index}')
scaler = getattr(self, f'scaler{index}')
scheduler = getattr(self, f'scheduler{index}')
warmup_scheduler = getattr(self, f'warmup{index}')
# set the grad scaler on the accelerator, since we are managing one per u-net
if exists(self.max_grad_norm):
self.accelerator.clip_grad_norm_(unet.parameters(), self.max_grad_norm)
optimizer.step()
optimizer.zero_grad()
if self.use_ema:
ema_unet = self.get_ema_unet(unet_number)
ema_unet.update()
# scheduler, if needed
maybe_warmup_context = nullcontext() if not exists(warmup_scheduler) else warmup_scheduler.dampening()
with maybe_warmup_context:
if exists(scheduler) and not self.accelerator.optimizer_step_was_skipped: # recommended in the docs
scheduler.step()
self.steps += F.one_hot(torch.tensor(unet_number - 1, device = self.steps.device), num_classes = len(self.steps))
if not exists(self.checkpoint_path):
return
total_steps = int(self.steps.sum().item())
if total_steps % self.checkpoint_every:
return
self.save_to_checkpoint_folder()
@torch.no_grad()
@cast_torch_tensor
@imagen_sample_in_chunks
def sample(self, *args, **kwargs):
context = nullcontext if kwargs.pop('use_non_ema', False) else self.use_ema_unets
self.print_untrained_unets()
if not self.is_main:
kwargs['use_tqdm'] = False
with context():
output = self.imagen.sample(*args, device = self.device, **kwargs)
return output
@partial(cast_torch_tensor, cast_fp16 = True)
def forward(
self,
*args,
unet_number = None,
max_batch_size = None,
**kwargs
):
unet_number = self.validate_unet_number(unet_number)
self.validate_and_set_unet_being_trained(unet_number)
self.set_accelerator_scaler(unet_number)
assert not exists(self.only_train_unet_number) or self.only_train_unet_number == unet_number, f'you can only train unet #{self.only_train_unet_number}'
total_loss = 0.
# + for debug
if self.CKeys['Debug_TrainerPack']==1:
print("In Trainer:Forward, check inputs:")
print('args: ', len(args))
print('args in:',args[0].shape)
print('kwargs: ', kwargs.keys())
for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs):
# + for debug
if self.CKeys['Debug_TrainerPack']==1:
print("after chunks,...")
print('chun_frac: ', chunk_size_frac)
print('chun_args: ', chunked_args)
print('chun_kwargs: ', chunked_kwargs)
with self.accelerator.autocast():
loss = self.model(
*chunked_args,
unet_number = unet_number,
**chunked_kwargs
)
loss = loss * chunk_size_frac
# + for debug
if self.CKeys['Debug_TrainerPack']==1:
print('part chun loss: ', loss)
total_loss += loss#.item()
if self.training:
self.accelerator.backward(loss)
return total_loss
# ========================================================
#
class ImagenTrainer_ModelB(nn.Module):
locked = False
def __init__(
self,
#imagen = None,
model = None,
imagen_checkpoint_path = None,
use_ema = True,
lr = 1e-4,
eps = 1e-8,
beta1 = 0.9,
beta2 = 0.99,
max_grad_norm = None,
group_wd_params = True,
warmup_steps = None,
cosine_decay_max_steps = None,
only_train_unet_number = None,
fp16 = False,
precision = None,
split_batches = True,
dl_tuple_output_keywords_names = ('images', 'text_embeds', 'text_masks', 'cond_images'),
verbose = True,
split_valid_fraction = 0.025,
split_valid_from_train = False,
split_random_seed = 42,
checkpoint_path = None,
checkpoint_every = None,
checkpoint_fs = None,
fs_kwargs: dict = None,
max_checkpoints_keep = 20,
# +++++++++++++++++++++
CKeys=None,
#
**kwargs
):
super().__init__()
assert not ImagenTrainer.locked, 'ImagenTrainer can only be initialized once per process - for the sake of distributed training, you will now have to create a separate script to train each unet (or a script that accepts unet number as an argument)'
assert exists(model.imagen) ^ exists(imagen_checkpoint_path), 'either imagen instance is passed into the trainer, or a checkpoint path that contains the imagen config'
# determine filesystem, using fsspec, for saving to local filesystem or cloud
self.fs = checkpoint_fs
if not exists(self.fs):
fs_kwargs = default(fs_kwargs, {})
self.fs, _ = url_to_fs(default(checkpoint_path, './'), **fs_kwargs)
# # -----------------------------------
# # from MJB
# assert isinstance(model.imagen, (ProteinDesigner_B))
# modified by BN
# ++
assert isinstance(model, (ProteinDesigner_B))
# +++++++++++++++++++++++++
self.CKeys = CKeys
ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)
self.imagen = model.imagen
self.model=model
self.is_elucidated = self.model.is_elucidated
# create accelerator instance
accelerate_kwargs, kwargs = groupby_prefix_and_trim('accelerate_', kwargs)
assert not (fp16 and exists(precision)), 'either set fp16 = True or forward the precision ("fp16", "bf16") to Accelerator'
accelerator_mixed_precision = default(precision, 'fp16' if fp16 else 'no')
self.accelerator = Accelerator(**{
'split_batches': split_batches,
'mixed_precision': accelerator_mixed_precision,
'kwargs_handlers': [DistributedDataParallelKwargs(find_unused_parameters = True)]
, **accelerate_kwargs})
ImagenTrainer.locked = self.is_distributed
# cast data to fp16 at training time if needed
self.cast_half_at_training = accelerator_mixed_precision == 'fp16'
# grad scaler must be managed outside of accelerator
grad_scaler_enabled = fp16
self.num_unets = len(self.imagen.unets)
self.use_ema = use_ema and self.is_main
self.ema_unets = nn.ModuleList([])
# keep track of what unet is being trained on
# only going to allow 1 unet training at a time
self.ema_unet_being_trained_index = -1 # keeps track of which ema unet is being trained on
# data related functions
self.train_dl_iter = None
self.train_dl = None
self.valid_dl_iter = None
self.valid_dl = None
self.dl_tuple_output_keywords_names = dl_tuple_output_keywords_names
# auto splitting validation from training, if dataset is passed in
self.split_valid_from_train = split_valid_from_train
assert 0 <= split_valid_fraction <= 1, 'split valid fraction must be between 0 and 1'
self.split_valid_fraction = split_valid_fraction
self.split_random_seed = split_random_seed
# be able to finely customize learning rate, weight decay
# per unet
lr, eps, warmup_steps, cosine_decay_max_steps = map(partial(cast_tuple, length = self.num_unets), (lr, eps, warmup_steps, cosine_decay_max_steps))
for ind, (unet, unet_lr, unet_eps, unet_warmup_steps, unet_cosine_decay_max_steps) in enumerate(zip(self.imagen.unets, lr, eps, warmup_steps, cosine_decay_max_steps)):
optimizer = Adam(
unet.parameters(),
lr = unet_lr,
eps = unet_eps,
betas = (beta1, beta2),
**kwargs
)
if self.use_ema:
self.ema_unets.append(EMA(unet, **ema_kwargs))
scaler = GradScaler(enabled = grad_scaler_enabled)
scheduler = warmup_scheduler = None
if exists(unet_cosine_decay_max_steps):
scheduler = CosineAnnealingLR(optimizer, T_max = unet_cosine_decay_max_steps)
if exists(unet_warmup_steps):
warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period = unet_warmup_steps)
if not exists(scheduler):
scheduler = LambdaLR(optimizer, lr_lambda = lambda step: 1.0)
# set on object
setattr(self, f'optim{ind}', optimizer) # cannot use pytorch ModuleList for some reason with optimizers
setattr(self, f'scaler{ind}', scaler)
setattr(self, f'scheduler{ind}', scheduler)
setattr(self, f'warmup{ind}', warmup_scheduler)
# gradient clipping if needed
self.max_grad_norm = max_grad_norm
# step tracker and misc
self.register_buffer('steps', torch.tensor([0] * self.num_unets))
self.verbose = verbose
# automatic set devices based on what accelerator decided
self.imagen.to(self.device)
self.to(self.device)
# checkpointing
assert not (exists(checkpoint_path) ^ exists(checkpoint_every))
self.checkpoint_path = checkpoint_path
self.checkpoint_every = checkpoint_every
self.max_checkpoints_keep = max_checkpoints_keep
self.can_checkpoint = self.is_local_main if isinstance(checkpoint_fs, LocalFileSystem) else self.is_main
if exists(checkpoint_path) and self.can_checkpoint:
bucket = url_to_bucket(checkpoint_path)
if not self.fs.exists(bucket):
self.fs.mkdir(bucket)
self.load_from_checkpoint_folder()
# only allowing training for unet
self.only_train_unet_number = only_train_unet_number
self.validate_and_set_unet_being_trained(only_train_unet_number)
# computed values
@property
def device(self):
return self.accelerator.device
@property
def is_distributed(self):
return not (self.accelerator.distributed_type == DistributedType.NO and self.accelerator.num_processes == 1)
@property
def is_main(self):
return self.accelerator.is_main_process
@property
def is_local_main(self):
return self.accelerator.is_local_main_process
@property
def unwrapped_unet(self):
return self.accelerator.unwrap_model(self.unet_being_trained)
# optimizer helper functions
def get_lr(self, unet_number):
self.validate_unet_number(unet_number)
unet_index = unet_number - 1
optim = getattr(self, f'optim{unet_index}')
return optim.param_groups[0]['lr']
# function for allowing only one unet from being trained at a time
def validate_and_set_unet_being_trained(self, unet_number = None):
if exists(unet_number):
self.validate_unet_number(unet_number)
assert not exists(self.only_train_unet_number) or self.only_train_unet_number == unet_number, 'you cannot only train on one unet at a time. you will need to save the trainer into a checkpoint, and resume training on a new unet'
self.only_train_unet_number = unet_number
self.imagen.only_train_unet_number = unet_number
if not exists(unet_number):
return
self.wrap_unet(unet_number)
def wrap_unet(self, unet_number):
if hasattr(self, 'one_unet_wrapped'):
return
unet = self.imagen.get_unet(unet_number)
self.unet_being_trained = self.accelerator.prepare(unet)
unet_index = unet_number - 1
optimizer = getattr(self, f'optim{unet_index}')
scheduler = getattr(self, f'scheduler{unet_index}')
optimizer = self.accelerator.prepare(optimizer)
if exists(scheduler):
scheduler = self.accelerator.prepare(scheduler)
setattr(self, f'optim{unet_index}', optimizer)
setattr(self, f'scheduler{unet_index}', scheduler)
self.one_unet_wrapped = True
# hacking accelerator due to not having separate gradscaler per optimizer
def set_accelerator_scaler(self, unet_number):
unet_number = self.validate_unet_number(unet_number)
scaler = getattr(self, f'scaler{unet_number - 1}')
self.accelerator.scaler = scaler
for optimizer in self.accelerator._optimizers:
optimizer.scaler = scaler
# helper print
def print(self, msg):
if not self.is_main:
return
if not self.verbose:
return
return self.accelerator.print(msg)
# validating the unet number
def validate_unet_number(self, unet_number = None):
if self.num_unets == 1:
unet_number = default(unet_number, 1)
assert 0 < unet_number <= self.num_unets, f'unet number should be in between 1 and {self.num_unets}'
return unet_number
# number of training steps taken
def num_steps_taken(self, unet_number = None):
if self.num_unets == 1:
unet_number = default(unet_number, 1)
return self.steps[unet_number - 1].item()
def print_untrained_unets(self):
print_final_error = False
for ind, (steps, unet) in enumerate(zip(self.steps.tolist(), self.imagen.unets)):
if steps > 0 or isinstance(unet, NullUnet):
continue
self.print(f'unet {ind + 1} has not been trained')
print_final_error = True
if print_final_error:
self.print('when sampling, you can pass stop_at_unet_number to stop early in the cascade, so it does not try to generate with untrained unets')
# data related functions
def add_train_dataloader(self, dl = None):
if not exists(dl):
return
assert not exists(self.train_dl), 'training dataloader was already added'
self.train_dl = self.accelerator.prepare(dl)
def add_valid_dataloader(self, dl):
if not exists(dl):
return
assert not exists(self.valid_dl), 'validation dataloader was already added'
self.valid_dl = self.accelerator.prepare(dl)
def add_train_dataset(self, ds = None, *, batch_size, **dl_kwargs):
if not exists(ds):
return
assert not exists(self.train_dl), 'training dataloader was already added'
valid_ds = None
if self.split_valid_from_train:
train_size = int((1 - self.split_valid_fraction) * len(ds))
valid_size = len(ds) - train_size
ds, valid_ds = random_split(ds, [train_size, valid_size], generator = torch.Generator().manual_seed(self.split_random_seed))
self.print(f'training with dataset of {len(ds)} samples and validating with randomly splitted {len(valid_ds)} samples')
dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
self.train_dl = self.accelerator.prepare(dl)
if not self.split_valid_from_train:
return
self.add_valid_dataset(valid_ds, batch_size = batch_size, **dl_kwargs)
def add_valid_dataset(self, ds, *, batch_size, **dl_kwargs):
if not exists(ds):
return
assert not exists(self.valid_dl), 'validation dataloader was already added'
dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
self.valid_dl = self.accelerator.prepare(dl)
def create_train_iter(self):
assert exists(self.train_dl), 'training dataloader has not been registered with the trainer yet'
if exists(self.train_dl_iter):
return
self.train_dl_iter = cycle(self.train_dl)
def create_valid_iter(self):
assert exists(self.valid_dl), 'validation dataloader has not been registered with the trainer yet'
if exists(self.valid_dl_iter):
return
self.valid_dl_iter = cycle(self.valid_dl)
def train_step(self, unet_number = None, **kwargs):
self.create_train_iter()
loss = self.step_with_dl_iter(self.train_dl_iter, unet_number = unet_number, **kwargs)
self.update(unet_number = unet_number)
return loss
@torch.no_grad()
@eval_decorator
def valid_step(self, **kwargs):
self.create_valid_iter()
context = self.use_ema_unets if kwargs.pop('use_ema_unets', False) else nullcontext
with context():
loss = self.step_with_dl_iter(self.valid_dl_iter, **kwargs)
return loss
def step_with_dl_iter(self, dl_iter, **kwargs):
dl_tuple_output = cast_tuple(next(dl_iter))
model_input = dict(list(zip(self.dl_tuple_output_keywords_names, dl_tuple_output)))
loss = self.forward(**{**kwargs, **model_input})
return loss
# checkpointing functions
@property
def all_checkpoints_sorted(self):
glob_pattern = os.path.join(self.checkpoint_path, '*.pt')
checkpoints = self.fs.glob(glob_pattern)
sorted_checkpoints = sorted(checkpoints, key = lambda x: int(str(x).split('.')[-2]), reverse = True)
return sorted_checkpoints
def load_from_checkpoint_folder(self, last_total_steps = -1):
if last_total_steps != -1:
filepath = os.path.join(self.checkpoint_path, f'checkpoint.{last_total_steps}.pt')
self.load(filepath)
return
sorted_checkpoints = self.all_checkpoints_sorted
if len(sorted_checkpoints) == 0:
self.print(f'no checkpoints found to load from at {self.checkpoint_path}')
return
last_checkpoint = sorted_checkpoints[0]
self.load(last_checkpoint)
def save_to_checkpoint_folder(self):
self.accelerator.wait_for_everyone()
if not self.can_checkpoint:
return
total_steps = int(self.steps.sum().item())
filepath = os.path.join(self.checkpoint_path, f'checkpoint.{total_steps}.pt')
self.save(filepath)
if self.max_checkpoints_keep <= 0:
return
sorted_checkpoints = self.all_checkpoints_sorted
checkpoints_to_discard = sorted_checkpoints[self.max_checkpoints_keep:]
for checkpoint in checkpoints_to_discard:
self.fs.rm(checkpoint)
# saving and loading functions
def save(
self,
path,
overwrite = True,
without_optim_and_sched = False,
**kwargs
):
self.accelerator.wait_for_everyone()
if not self.can_checkpoint:
return
fs = self.fs
assert not (fs.exists(path) and not overwrite)
self.reset_ema_unets_all_one_device()
save_obj = dict(
model = self.imagen.state_dict(),
version = __version__,
steps = self.steps.cpu(),
**kwargs
)
save_optim_and_sched_iter = range(0, self.num_unets) if not without_optim_and_sched else tuple()
for ind in save_optim_and_sched_iter:
scaler_key = f'scaler{ind}'
optimizer_key = f'optim{ind}'
scheduler_key = f'scheduler{ind}'
warmup_scheduler_key = f'warmup{ind}'
scaler = getattr(self, scaler_key)
optimizer = getattr(self, optimizer_key)
scheduler = getattr(self, scheduler_key)
warmup_scheduler = getattr(self, warmup_scheduler_key)
if exists(scheduler):
save_obj = {**save_obj, scheduler_key: scheduler.state_dict()}
if exists(warmup_scheduler):
save_obj = {**save_obj, warmup_scheduler_key: warmup_scheduler.state_dict()}
save_obj = {**save_obj, scaler_key: scaler.state_dict(), optimizer_key: optimizer.state_dict()}
if self.use_ema:
save_obj = {**save_obj, 'ema': self.ema_unets.state_dict()}
# determine if imagen config is available
if hasattr(self.imagen, '_config'):
self.print(f'this checkpoint is commandable from the CLI - "imagen --model {str(path)} \"<prompt>\""')
save_obj = {
**save_obj,
'imagen_type': 'elucidated' if self.is_elucidated else 'original',
'imagen_params': self.imagen._config
}
#save to path
with fs.open(path, 'wb') as f:
torch.save(save_obj, f)
self.print(f'checkpoint saved to {path}')
def load(self, path, only_model = False, strict = True, noop_if_not_exist = False):
fs = self.fs
if noop_if_not_exist and not fs.exists(path):
self.print(f'trainer checkpoint not found at {str(path)}')
return
assert fs.exists(path), f'{path} does not exist'
self.reset_ema_unets_all_one_device()
# to avoid extra GPU memory usage in main process when using Accelerate
with fs.open(path) as f:
loaded_obj = torch.load(f, map_location='cpu')
if version.parse(__version__) != version.parse(loaded_obj['version']):
self.print(f'loading saved imagen at version {loaded_obj["version"]}, but current package version is {__version__}')
try:
self.imagen.load_state_dict(loaded_obj['model'], strict = strict)
except RuntimeError:
print("Failed loading state dict. Trying partial load")
self.imagen.load_state_dict(restore_parts(self.imagen.state_dict(),
loaded_obj['model']))
if only_model:
return loaded_obj
self.steps.copy_(loaded_obj['steps'])
for ind in range(0, self.num_unets):
scaler_key = f'scaler{ind}'
optimizer_key = f'optim{ind}'
scheduler_key = f'scheduler{ind}'
warmup_scheduler_key = f'warmup{ind}'
scaler = getattr(self, scaler_key)
optimizer = getattr(self, optimizer_key)
scheduler = getattr(self, scheduler_key)
warmup_scheduler = getattr(self, warmup_scheduler_key)
if exists(scheduler) and scheduler_key in loaded_obj:
scheduler.load_state_dict(loaded_obj[scheduler_key])
if exists(warmup_scheduler) and warmup_scheduler_key in loaded_obj:
warmup_scheduler.load_state_dict(loaded_obj[warmup_scheduler_key])
if exists(optimizer):
try:
optimizer.load_state_dict(loaded_obj[optimizer_key])
scaler.load_state_dict(loaded_obj[scaler_key])
except:
self.print('could not load optimizer and scaler, possibly because you have turned on mixed precision training since the last run. resuming with new optimizer and scalers')
if self.use_ema:
assert 'ema' in loaded_obj
try:
self.ema_unets.load_state_dict(loaded_obj['ema'], strict = strict)
except RuntimeError:
print("Failed loading state dict. Trying partial load")
self.ema_unets.load_state_dict(restore_parts(self.ema_unets.state_dict(),
loaded_obj['ema']))
self.print(f'checkpoint loaded from {path}')
return loaded_obj
# managing ema unets and their devices
@property
def unets(self):
return nn.ModuleList([ema.ema_model for ema in self.ema_unets])
def get_ema_unet(self, unet_number = None):
if not self.use_ema:
return
unet_number = self.validate_unet_number(unet_number)
index = unet_number - 1
if isinstance(self.unets, nn.ModuleList):
unets_list = [unet for unet in self.ema_unets]
delattr(self, 'ema_unets')
self.ema_unets = unets_list
if index != self.ema_unet_being_trained_index:
for unet_index, unet in enumerate(self.ema_unets):
unet.to(self.device if unet_index == index else 'cpu')
self.ema_unet_being_trained_index = index
return self.ema_unets[index]
def reset_ema_unets_all_one_device(self, device = None):
if not self.use_ema:
return
device = default(device, self.device)
self.ema_unets = nn.ModuleList([*self.ema_unets])
self.ema_unets.to(device)
self.ema_unet_being_trained_index = -1
@torch.no_grad()
@contextmanager
def use_ema_unets(self):
if not self.use_ema:
output = yield
return output
self.reset_ema_unets_all_one_device()
self.imagen.reset_unets_all_one_device()
self.unets.eval()
trainable_unets = self.imagen.unets
self.imagen.unets = self.unets # swap in exponential moving averaged unets for sampling
output = yield
self.imagen.unets = trainable_unets # restore original training unets
# cast the ema_model unets back to original device
for ema in self.ema_unets:
ema.restore_ema_model_device()
return output
def print_unet_devices(self):
self.print('unet devices:')
for i, unet in enumerate(self.imagen.unets):
device = next(unet.parameters()).device
self.print(f'\tunet {i}: {device}')
if not self.use_ema:
return
self.print('\nema unet devices:')
for i, ema_unet in enumerate(self.ema_unets):
device = next(ema_unet.parameters()).device
self.print(f'\tema unet {i}: {device}')
# overriding state dict functions
def state_dict(self, *args, **kwargs):
self.reset_ema_unets_all_one_device()
return super().state_dict(*args, **kwargs)
def load_state_dict(self, *args, **kwargs):
self.reset_ema_unets_all_one_device()
return super().load_state_dict(*args, **kwargs)
# encoding text functions
def encode_text(self, text, **kwargs):
return self.imagen.encode_text(text, **kwargs)
# forwarding functions and gradient step updates
def update(self, unet_number = None):
unet_number = self.validate_unet_number(unet_number)
self.validate_and_set_unet_being_trained(unet_number)
self.set_accelerator_scaler(unet_number)
index = unet_number - 1
unet = self.unet_being_trained
optimizer = getattr(self, f'optim{index}')
scaler = getattr(self, f'scaler{index}')
scheduler = getattr(self, f'scheduler{index}')
warmup_scheduler = getattr(self, f'warmup{index}')
# set the grad scaler on the accelerator, since we are managing one per u-net
if exists(self.max_grad_norm):
self.accelerator.clip_grad_norm_(unet.parameters(), self.max_grad_norm)
optimizer.step()
optimizer.zero_grad()
if self.use_ema:
ema_unet = self.get_ema_unet(unet_number)
ema_unet.update()
# scheduler, if needed
maybe_warmup_context = nullcontext() if not exists(warmup_scheduler) else warmup_scheduler.dampening()
with maybe_warmup_context:
if exists(scheduler) and not self.accelerator.optimizer_step_was_skipped: # recommended in the docs
scheduler.step()
self.steps += F.one_hot(torch.tensor(unet_number - 1, device = self.steps.device), num_classes = len(self.steps))
if not exists(self.checkpoint_path):
return
total_steps = int(self.steps.sum().item())
if total_steps % self.checkpoint_every:
return
self.save_to_checkpoint_folder()
@torch.no_grad()
@cast_torch_tensor
@imagen_sample_in_chunks
def sample(self, *args, **kwargs):
context = nullcontext if kwargs.pop('use_non_ema', False) else self.use_ema_unets
self.print_untrained_unets()
if not self.is_main:
kwargs['use_tqdm'] = False
with context():
output = self.imagen.sample(*args, device = self.device, **kwargs)
return output
@partial(cast_torch_tensor, cast_fp16 = True)
def forward(
self,
*args,
unet_number = None,
max_batch_size = None,
**kwargs
):
unet_number = self.validate_unet_number(unet_number)
self.validate_and_set_unet_being_trained(unet_number)
self.set_accelerator_scaler(unet_number)
assert not exists(self.only_train_unet_number) or self.only_train_unet_number == unet_number, f'you can only train unet #{self.only_train_unet_number}'
total_loss = 0.
# + for debug
if self.CKeys['Debug_TrainerPack']==1:
print("In Trainer:Forward, check inputs:")
print('args: ', len(args))
print('args in:',args[0].shape)
print('kwargs: ', kwargs.keys())
for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs):
# + for debug
if self.CKeys['Debug_TrainerPack']==1:
print("after chunks,...")
print('chun_frac: ', chunk_size_frac)
print('chun_args: ', chunked_args)
print('chun_kwargs: ', chunked_kwargs)
with self.accelerator.autocast():
loss = self.model(
*chunked_args,
unet_number = unet_number,
**chunked_kwargs
)
loss = loss * chunk_size_frac
# + for debug
if self.CKeys['Debug_TrainerPack']==1:
print('part chun loss: ', loss)
total_loss += loss#.item()
if self.training:
self.accelerator.backward(loss)
return total_loss
class ImagenTrainer_Old(nn.Module):
locked = False
def __init__(
self,
#imagen = None,
model = None,
imagen_checkpoint_path = None,
use_ema = True,
lr = 1e-4,
eps = 1e-8,
beta1 = 0.9,
beta2 = 0.99,
max_grad_norm = None,
group_wd_params = True,
warmup_steps = None,
cosine_decay_max_steps = None,
only_train_unet_number = None,
fp16 = False,
precision = None,
split_batches = True,
dl_tuple_output_keywords_names = ('images', 'text_embeds', 'text_masks', 'cond_images'),
verbose = True,
split_valid_fraction = 0.025,
split_valid_from_train = False,
split_random_seed = 42,
checkpoint_path = None,
checkpoint_every = None,
checkpoint_fs = None,
fs_kwargs: dict = None,
max_checkpoints_keep = 20,
**kwargs
):
super().__init__()
assert not ImagenTrainer.locked, 'ImagenTrainer can only be initialized once per process - for the sake of distributed training, you will now have to create a separate script to train each unet (or a script that accepts unet number as an argument)'
assert exists(model.imagen) ^ exists(imagen_checkpoint_path), 'either imagen instance is passed into the trainer, or a checkpoint path that contains the imagen config'
# determine filesystem, using fsspec, for saving to local filesystem or cloud
self.fs = checkpoint_fs
if not exists(self.fs):
fs_kwargs = default(fs_kwargs, {})
self.fs, _ = url_to_fs(default(checkpoint_path, './'), **fs_kwargs)
# # -----------------------------------
# # from MJB
# assert isinstance(model.imagen, (ProteinDesigner_B))
ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)
self.imagen = model.imagen
self.model=model
self.is_elucidated = self.model.is_elucidated
# create accelerator instance
accelerate_kwargs, kwargs = groupby_prefix_and_trim('accelerate_', kwargs)
assert not (fp16 and exists(precision)), 'either set fp16 = True or forward the precision ("fp16", "bf16") to Accelerator'
accelerator_mixed_precision = default(precision, 'fp16' if fp16 else 'no')
self.accelerator = Accelerator(**{
'split_batches': split_batches,
'mixed_precision': accelerator_mixed_precision,
'kwargs_handlers': [DistributedDataParallelKwargs(find_unused_parameters = True)]
, **accelerate_kwargs})
ImagenTrainer.locked = self.is_distributed
# cast data to fp16 at training time if needed
self.cast_half_at_training = accelerator_mixed_precision == 'fp16'
# grad scaler must be managed outside of accelerator
grad_scaler_enabled = fp16
self.num_unets = len(self.imagen.unets)
self.use_ema = use_ema and self.is_main
self.ema_unets = nn.ModuleList([])
# keep track of what unet is being trained on
# only going to allow 1 unet training at a time
self.ema_unet_being_trained_index = -1 # keeps track of which ema unet is being trained on
# data related functions
self.train_dl_iter = None
self.train_dl = None
self.valid_dl_iter = None
self.valid_dl = None
self.dl_tuple_output_keywords_names = dl_tuple_output_keywords_names
# auto splitting validation from training, if dataset is passed in
self.split_valid_from_train = split_valid_from_train
assert 0 <= split_valid_fraction <= 1, 'split valid fraction must be between 0 and 1'
self.split_valid_fraction = split_valid_fraction
self.split_random_seed = split_random_seed
# be able to finely customize learning rate, weight decay
# per unet
lr, eps, warmup_steps, cosine_decay_max_steps = map(partial(cast_tuple, length = self.num_unets), (lr, eps, warmup_steps, cosine_decay_max_steps))
for ind, (unet, unet_lr, unet_eps, unet_warmup_steps, unet_cosine_decay_max_steps) in enumerate(zip(self.imagen.unets, lr, eps, warmup_steps, cosine_decay_max_steps)):
optimizer = Adam(
unet.parameters(),
lr = unet_lr,
eps = unet_eps,
betas = (beta1, beta2),
**kwargs
)
if self.use_ema:
self.ema_unets.append(EMA(unet, **ema_kwargs))
scaler = GradScaler(enabled = grad_scaler_enabled)
scheduler = warmup_scheduler = None
if exists(unet_cosine_decay_max_steps):
scheduler = CosineAnnealingLR(optimizer, T_max = unet_cosine_decay_max_steps)
if exists(unet_warmup_steps):
warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period = unet_warmup_steps)
if not exists(scheduler):
scheduler = LambdaLR(optimizer, lr_lambda = lambda step: 1.0)
# set on object
setattr(self, f'optim{ind}', optimizer) # cannot use pytorch ModuleList for some reason with optimizers
setattr(self, f'scaler{ind}', scaler)
setattr(self, f'scheduler{ind}', scheduler)
setattr(self, f'warmup{ind}', warmup_scheduler)
# gradient clipping if needed
self.max_grad_norm = max_grad_norm
# step tracker and misc
self.register_buffer('steps', torch.tensor([0] * self.num_unets))
self.verbose = verbose
# automatic set devices based on what accelerator decided
self.imagen.to(self.device)
self.to(self.device)
# checkpointing
assert not (exists(checkpoint_path) ^ exists(checkpoint_every))
self.checkpoint_path = checkpoint_path
self.checkpoint_every = checkpoint_every
self.max_checkpoints_keep = max_checkpoints_keep
self.can_checkpoint = self.is_local_main if isinstance(checkpoint_fs, LocalFileSystem) else self.is_main
if exists(checkpoint_path) and self.can_checkpoint:
bucket = url_to_bucket(checkpoint_path)
if not self.fs.exists(bucket):
self.fs.mkdir(bucket)
self.load_from_checkpoint_folder()
# only allowing training for unet
self.only_train_unet_number = only_train_unet_number
self.validate_and_set_unet_being_trained(only_train_unet_number)
# computed values
@property
def device(self):
return self.accelerator.device
@property
def is_distributed(self):
return not (self.accelerator.distributed_type == DistributedType.NO and self.accelerator.num_processes == 1)
@property
def is_main(self):
return self.accelerator.is_main_process
@property
def is_local_main(self):
return self.accelerator.is_local_main_process
@property
def unwrapped_unet(self):
return self.accelerator.unwrap_model(self.unet_being_trained)
# optimizer helper functions
def get_lr(self, unet_number):
self.validate_unet_number(unet_number)
unet_index = unet_number - 1
optim = getattr(self, f'optim{unet_index}')
return optim.param_groups[0]['lr']
# function for allowing only one unet from being trained at a time
def validate_and_set_unet_being_trained(self, unet_number = None):
if exists(unet_number):
self.validate_unet_number(unet_number)
assert not exists(self.only_train_unet_number) or self.only_train_unet_number == unet_number, 'you cannot only train on one unet at a time. you will need to save the trainer into a checkpoint, and resume training on a new unet'
self.only_train_unet_number = unet_number
self.imagen.only_train_unet_number = unet_number
if not exists(unet_number):
return
self.wrap_unet(unet_number)
def wrap_unet(self, unet_number):
if hasattr(self, 'one_unet_wrapped'):
return
unet = self.imagen.get_unet(unet_number)
self.unet_being_trained = self.accelerator.prepare(unet)
unet_index = unet_number - 1
optimizer = getattr(self, f'optim{unet_index}')
scheduler = getattr(self, f'scheduler{unet_index}')
optimizer = self.accelerator.prepare(optimizer)
if exists(scheduler):
scheduler = self.accelerator.prepare(scheduler)
setattr(self, f'optim{unet_index}', optimizer)
setattr(self, f'scheduler{unet_index}', scheduler)
self.one_unet_wrapped = True
# hacking accelerator due to not having separate gradscaler per optimizer
def set_accelerator_scaler(self, unet_number):
unet_number = self.validate_unet_number(unet_number)
scaler = getattr(self, f'scaler{unet_number - 1}')
self.accelerator.scaler = scaler
for optimizer in self.accelerator._optimizers:
optimizer.scaler = scaler
# helper print
def print(self, msg):
if not self.is_main:
return
if not self.verbose:
return
return self.accelerator.print(msg)
# validating the unet number
def validate_unet_number(self, unet_number = None):
if self.num_unets == 1:
unet_number = default(unet_number, 1)
assert 0 < unet_number <= self.num_unets, f'unet number should be in between 1 and {self.num_unets}'
return unet_number
# number of training steps taken
def num_steps_taken(self, unet_number = None):
if self.num_unets == 1:
unet_number = default(unet_number, 1)
return self.steps[unet_number - 1].item()
def print_untrained_unets(self):
print_final_error = False
for ind, (steps, unet) in enumerate(zip(self.steps.tolist(), self.imagen.unets)):
if steps > 0 or isinstance(unet, NullUnet):
continue
self.print(f'unet {ind + 1} has not been trained')
print_final_error = True
if print_final_error:
self.print('when sampling, you can pass stop_at_unet_number to stop early in the cascade, so it does not try to generate with untrained unets')
# data related functions
def add_train_dataloader(self, dl = None):
if not exists(dl):
return
assert not exists(self.train_dl), 'training dataloader was already added'
self.train_dl = self.accelerator.prepare(dl)
def add_valid_dataloader(self, dl):
if not exists(dl):
return
assert not exists(self.valid_dl), 'validation dataloader was already added'
self.valid_dl = self.accelerator.prepare(dl)
def add_train_dataset(self, ds = None, *, batch_size, **dl_kwargs):
if not exists(ds):
return
assert not exists(self.train_dl), 'training dataloader was already added'
valid_ds = None
if self.split_valid_from_train:
train_size = int((1 - self.split_valid_fraction) * len(ds))
valid_size = len(ds) - train_size
ds, valid_ds = random_split(ds, [train_size, valid_size], generator = torch.Generator().manual_seed(self.split_random_seed))
self.print(f'training with dataset of {len(ds)} samples and validating with randomly splitted {len(valid_ds)} samples')
dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
self.train_dl = self.accelerator.prepare(dl)
if not self.split_valid_from_train:
return
self.add_valid_dataset(valid_ds, batch_size = batch_size, **dl_kwargs)
def add_valid_dataset(self, ds, *, batch_size, **dl_kwargs):
if not exists(ds):
return
assert not exists(self.valid_dl), 'validation dataloader was already added'
dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
self.valid_dl = self.accelerator.prepare(dl)
def create_train_iter(self):
assert exists(self.train_dl), 'training dataloader has not been registered with the trainer yet'
if exists(self.train_dl_iter):
return
self.train_dl_iter = cycle(self.train_dl)
def create_valid_iter(self):
assert exists(self.valid_dl), 'validation dataloader has not been registered with the trainer yet'
if exists(self.valid_dl_iter):
return
self.valid_dl_iter = cycle(self.valid_dl)
def train_step(self, unet_number = None, **kwargs):
self.create_train_iter()
loss = self.step_with_dl_iter(self.train_dl_iter, unet_number = unet_number, **kwargs)
self.update(unet_number = unet_number)
return loss
@torch.no_grad()
@eval_decorator
def valid_step(self, **kwargs):
self.create_valid_iter()
context = self.use_ema_unets if kwargs.pop('use_ema_unets', False) else nullcontext
with context():
loss = self.step_with_dl_iter(self.valid_dl_iter, **kwargs)
return loss
def step_with_dl_iter(self, dl_iter, **kwargs):
dl_tuple_output = cast_tuple(next(dl_iter))
model_input = dict(list(zip(self.dl_tuple_output_keywords_names, dl_tuple_output)))
loss = self.forward(**{**kwargs, **model_input})
return loss
# checkpointing functions
@property
def all_checkpoints_sorted(self):
glob_pattern = os.path.join(self.checkpoint_path, '*.pt')
checkpoints = self.fs.glob(glob_pattern)
sorted_checkpoints = sorted(checkpoints, key = lambda x: int(str(x).split('.')[-2]), reverse = True)
return sorted_checkpoints
def load_from_checkpoint_folder(self, last_total_steps = -1):
if last_total_steps != -1:
filepath = os.path.join(self.checkpoint_path, f'checkpoint.{last_total_steps}.pt')
self.load(filepath)
return
sorted_checkpoints = self.all_checkpoints_sorted
if len(sorted_checkpoints) == 0:
self.print(f'no checkpoints found to load from at {self.checkpoint_path}')
return
last_checkpoint = sorted_checkpoints[0]
self.load(last_checkpoint)
def save_to_checkpoint_folder(self):
self.accelerator.wait_for_everyone()
if not self.can_checkpoint:
return
total_steps = int(self.steps.sum().item())
filepath = os.path.join(self.checkpoint_path, f'checkpoint.{total_steps}.pt')
self.save(filepath)
if self.max_checkpoints_keep <= 0:
return
sorted_checkpoints = self.all_checkpoints_sorted
checkpoints_to_discard = sorted_checkpoints[self.max_checkpoints_keep:]
for checkpoint in checkpoints_to_discard:
self.fs.rm(checkpoint)
# saving and loading functions
def save(
self,
path,
overwrite = True,
without_optim_and_sched = False,
**kwargs
):
self.accelerator.wait_for_everyone()
if not self.can_checkpoint:
return
fs = self.fs
assert not (fs.exists(path) and not overwrite)
self.reset_ema_unets_all_one_device()
save_obj = dict(
model = self.imagen.state_dict(),
version = __version__,
steps = self.steps.cpu(),
**kwargs
)
save_optim_and_sched_iter = range(0, self.num_unets) if not without_optim_and_sched else tuple()
for ind in save_optim_and_sched_iter:
scaler_key = f'scaler{ind}'
optimizer_key = f'optim{ind}'
scheduler_key = f'scheduler{ind}'
warmup_scheduler_key = f'warmup{ind}'
scaler = getattr(self, scaler_key)
optimizer = getattr(self, optimizer_key)
scheduler = getattr(self, scheduler_key)
warmup_scheduler = getattr(self, warmup_scheduler_key)
if exists(scheduler):
save_obj = {**save_obj, scheduler_key: scheduler.state_dict()}
if exists(warmup_scheduler):
save_obj = {**save_obj, warmup_scheduler_key: warmup_scheduler.state_dict()}
save_obj = {**save_obj, scaler_key: scaler.state_dict(), optimizer_key: optimizer.state_dict()}
if self.use_ema:
save_obj = {**save_obj, 'ema': self.ema_unets.state_dict()}
# determine if imagen config is available
if hasattr(self.imagen, '_config'):
self.print(f'this checkpoint is commandable from the CLI - "imagen --model {str(path)} \"<prompt>\""')
save_obj = {
**save_obj,
'imagen_type': 'elucidated' if self.is_elucidated else 'original',
'imagen_params': self.imagen._config
}
#save to path
with fs.open(path, 'wb') as f:
torch.save(save_obj, f)
self.print(f'checkpoint saved to {path}')
def load(self, path, only_model = False, strict = True, noop_if_not_exist = False):
fs = self.fs
if noop_if_not_exist and not fs.exists(path):
self.print(f'trainer checkpoint not found at {str(path)}')
return
assert fs.exists(path), f'{path} does not exist'
self.reset_ema_unets_all_one_device()
# to avoid extra GPU memory usage in main process when using Accelerate
with fs.open(path) as f:
loaded_obj = torch.load(f, map_location='cpu')
if version.parse(__version__) != version.parse(loaded_obj['version']):
self.print(f'loading saved imagen at version {loaded_obj["version"]}, but current package version is {__version__}')
try:
self.imagen.load_state_dict(loaded_obj['model'], strict = strict)
except RuntimeError:
print("Failed loading state dict. Trying partial load")
self.imagen.load_state_dict(restore_parts(self.imagen.state_dict(),
loaded_obj['model']))
if only_model:
return loaded_obj
self.steps.copy_(loaded_obj['steps'])
for ind in range(0, self.num_unets):
scaler_key = f'scaler{ind}'
optimizer_key = f'optim{ind}'
scheduler_key = f'scheduler{ind}'
warmup_scheduler_key = f'warmup{ind}'
scaler = getattr(self, scaler_key)
optimizer = getattr(self, optimizer_key)
scheduler = getattr(self, scheduler_key)
warmup_scheduler = getattr(self, warmup_scheduler_key)
if exists(scheduler) and scheduler_key in loaded_obj:
scheduler.load_state_dict(loaded_obj[scheduler_key])
if exists(warmup_scheduler) and warmup_scheduler_key in loaded_obj:
warmup_scheduler.load_state_dict(loaded_obj[warmup_scheduler_key])
if exists(optimizer):
try:
optimizer.load_state_dict(loaded_obj[optimizer_key])
scaler.load_state_dict(loaded_obj[scaler_key])
except:
self.print('could not load optimizer and scaler, possibly because you have turned on mixed precision training since the last run. resuming with new optimizer and scalers')
if self.use_ema:
assert 'ema' in loaded_obj
try:
self.ema_unets.load_state_dict(loaded_obj['ema'], strict = strict)
except RuntimeError:
print("Failed loading state dict. Trying partial load")
self.ema_unets.load_state_dict(restore_parts(self.ema_unets.state_dict(),
loaded_obj['ema']))
self.print(f'checkpoint loaded from {path}')
return loaded_obj
# managing ema unets and their devices
@property
def unets(self):
return nn.ModuleList([ema.ema_model for ema in self.ema_unets])
def get_ema_unet(self, unet_number = None):
if not self.use_ema:
return
unet_number = self.validate_unet_number(unet_number)
index = unet_number - 1
if isinstance(self.unets, nn.ModuleList):
unets_list = [unet for unet in self.ema_unets]
delattr(self, 'ema_unets')
self.ema_unets = unets_list
if index != self.ema_unet_being_trained_index:
for unet_index, unet in enumerate(self.ema_unets):
unet.to(self.device if unet_index == index else 'cpu')
self.ema_unet_being_trained_index = index
return self.ema_unets[index]
def reset_ema_unets_all_one_device(self, device = None):
if not self.use_ema:
return
device = default(device, self.device)
self.ema_unets = nn.ModuleList([*self.ema_unets])
self.ema_unets.to(device)
self.ema_unet_being_trained_index = -1
@torch.no_grad()
@contextmanager
def use_ema_unets(self):
if not self.use_ema:
output = yield
return output
self.reset_ema_unets_all_one_device()
self.imagen.reset_unets_all_one_device()
self.unets.eval()
trainable_unets = self.imagen.unets
self.imagen.unets = self.unets # swap in exponential moving averaged unets for sampling
output = yield
self.imagen.unets = trainable_unets # restore original training unets
# cast the ema_model unets back to original device
for ema in self.ema_unets:
ema.restore_ema_model_device()
return output
def print_unet_devices(self):
self.print('unet devices:')
for i, unet in enumerate(self.imagen.unets):
device = next(unet.parameters()).device
self.print(f'\tunet {i}: {device}')
if not self.use_ema:
return
self.print('\nema unet devices:')
for i, ema_unet in enumerate(self.ema_unets):
device = next(ema_unet.parameters()).device
self.print(f'\tema unet {i}: {device}')
# overriding state dict functions
def state_dict(self, *args, **kwargs):
self.reset_ema_unets_all_one_device()
return super().state_dict(*args, **kwargs)
def load_state_dict(self, *args, **kwargs):
self.reset_ema_unets_all_one_device()
return super().load_state_dict(*args, **kwargs)
# encoding text functions
def encode_text(self, text, **kwargs):
return self.imagen.encode_text(text, **kwargs)
# forwarding functions and gradient step updates
def update(self, unet_number = None):
unet_number = self.validate_unet_number(unet_number)
self.validate_and_set_unet_being_trained(unet_number)
self.set_accelerator_scaler(unet_number)
index = unet_number - 1
unet = self.unet_being_trained
optimizer = getattr(self, f'optim{index}')
scaler = getattr(self, f'scaler{index}')
scheduler = getattr(self, f'scheduler{index}')
warmup_scheduler = getattr(self, f'warmup{index}')
# set the grad scaler on the accelerator, since we are managing one per u-net
if exists(self.max_grad_norm):
self.accelerator.clip_grad_norm_(unet.parameters(), self.max_grad_norm)
optimizer.step()
optimizer.zero_grad()
if self.use_ema:
ema_unet = self.get_ema_unet(unet_number)
ema_unet.update()
# scheduler, if needed
maybe_warmup_context = nullcontext() if not exists(warmup_scheduler) else warmup_scheduler.dampening()
with maybe_warmup_context:
if exists(scheduler) and not self.accelerator.optimizer_step_was_skipped: # recommended in the docs
scheduler.step()
self.steps += F.one_hot(torch.tensor(unet_number - 1, device = self.steps.device), num_classes = len(self.steps))
if not exists(self.checkpoint_path):
return
total_steps = int(self.steps.sum().item())
if total_steps % self.checkpoint_every:
return
self.save_to_checkpoint_folder()
@torch.no_grad()
@cast_torch_tensor
@imagen_sample_in_chunks
def sample(self, *args, **kwargs):
context = nullcontext if kwargs.pop('use_non_ema', False) else self.use_ema_unets
self.print_untrained_unets()
if not self.is_main:
kwargs['use_tqdm'] = False
with context():
output = self.imagen.sample(*args, device = self.device, **kwargs)
return output
@partial(cast_torch_tensor, cast_fp16 = True)
def forward(
self,
*args,
unet_number = None,
max_batch_size = None,
**kwargs
):
unet_number = self.validate_unet_number(unet_number)
self.validate_and_set_unet_being_trained(unet_number)
self.set_accelerator_scaler(unet_number)
assert not exists(self.only_train_unet_number) or self.only_train_unet_number == unet_number, f'you can only train unet #{self.only_train_unet_number}'
total_loss = 0.
# + for debug
print('args: ', len(args))
print('args in:',args[0].shape)
print('kwargs: ', kwargs.keys())
for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs):
# + for debug
print('chun_frac: ', chunk_size_frac)
print('chun_args: ', chunked_args)
print('chun_kwargs: ', chunked_kwargs)
with self.accelerator.autocast():
loss = self.model(*chunked_args, unet_number = unet_number, **chunked_kwargs)
loss = loss * chunk_size_frac
print('loss: ', loss)
total_loss += loss#.item()
if self.training:
self.accelerator.backward(loss)
print('I am here')
return total_loss
def write_fasta (sequence, filename_out):
with open (filename_out, mode ='w') as f:
f.write (f'>{filename_out}\n')
f.write (f'{sequence}')
#
def sample_sequence_FromModelB (
model,
X=None, #this is the target conventionally when using text embd
flag=0,
cond_scales=1.,
foldproteins=False,
X_string=None,
x_data=None,
skip_steps=0,
inpaint_images = None,
inpaint_masks = None,
inpaint_resample_times = None,
init_images = None,
num_cycle=16,
# ++++++++++++++++++++++++
ynormfac=1,
train_unet_number=1,
tokenizer_X=None,
Xnormfac=1.,
max_length=1.,
prefix=None,
tokenizer_y=None,
):
steps=0
e=flag
#num_samples = min (num_samples,y_train_batch.shape[0] )
if X!=None:
print (f"Producing {len(X)} samples...from text conditioning X...")
lenn_val=len(X)
if X_string!=None:
lenn_val=len(X_string)
print (f"Producing {len(X_string)} samples...from text conditioning X_String (from string)...")
if x_data!=None:
print (f"Producing {len(x_data)} samples...from image conditingig x_data ...")
lenn_val=len(x_data)
print (x_data)
print ('Device: ', device)
for iisample in range (lenn_val):
X_cond=None
if X_string==None and X != None: #only do if X provided
X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
if X_string !=None:
X = tokenizer_X.texts_to_sequences(X_string[iisample])
X= sequence.pad_sequences(X, maxlen=max_length, padding='post', truncating='post')
X=np.array(X)
X_cond=torch.from_numpy(X).float()/Xnormfac
print ('Tokenized and processed: ', X_cond)
print ("X_cond=", X_cond)
result=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales ,
x_data=x_data, skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,device=device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
)
result=torch.round(result*ynormfac)
plt.plot (result[0,0,:].cpu().detach().numpy(),label= f'Predicted')
#plt.plot (GT[samples,0,:]*ynormfac,label= f'GT {0}')
plt.legend()
outname = prefix+ f"sampled_from_X_{flag}_condscale-{str (cond_scales)}_{e}_{steps}.jpg"
#plt.title (f"Sample {samples}, cond scale={str (cond_scales[iisample])}")
plt.savefig(outname, dpi=200)
plt.show ()
to_rev=result[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
print (to_rev.shape)
y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
### reverse second structure input....
if X_cond != None:
X_cond=torch.round(X_cond*Xnormfac)
to_rev=X_cond[:,:]
to_rev=to_rev.long().cpu().detach().numpy()
print (to_rev.shape)
X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
if x_data !=None:
X_data_reversed=x_data #is already in sequence fromat..
print (f"For {X} or {X_data_reversed[iisample]}, predicted sequence: ", y_data_reversed[iisample])
if foldproteins:
if X_cond != None:
xbc=X_cond[iisample,:].cpu().detach().numpy()
out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})+f'_{flag}_{steps}'
if x_data !=None:
#xbc=x_data[iisample]
out_nam=x_data[iisample]
tempname='temp'
pdb_file=foldandsavePDB (
sequence=y_data_reversed[0],
filename_out=tempname,
num_cycle=num_cycle,
flag=flag,
# +++++++++++++++++++
prefix=prefix
)
out_nam_fasta=f'{prefix}{out_nam}_{flag}_{steps}.fasta'
write_fasta (y_data_reversed[0], out_nam_fasta)
out_nam=f'{prefix}{X_data_reversed[iisample]}_{flag}_{steps}.pdb'
# print('Debug 1: out: ', out_nam)
# print('Debug 2: in: ', pdb_file)
shutil.copy (pdb_file, out_nam) #source, dest
# cmd_line = 'cp ' + pdb_file + ' ' + out_nam
# print(cmd_line)
# os.popen(cmd_line)
# print('Debug 3')
pdb_file=out_nam
print (f"Properly named PDB file produced: {pdb_file}")
#flag=1000
view=show_pdb(pdb_file=pdb_file, flag=flag,
show_sidechains=show_sidechains, show_mainchains=show_mainchains, color=color)
view.show()
steps=steps+1
return pdb_file
def sample_loop_FromModelB (model,
train_loader,
cond_scales=[7.5], #list of cond scales - each sampled...
num_samples=2, #how many samples produced every time tested.....
timesteps=100,
flag=0,foldproteins=False,
use_text_embedd=True,skip_steps=0,
# +++++++++++++++++++
train_unet_number=1,
ynormfac=1,
prefix=None,
tokenizer_y=None,
Xnormfac=1,
tokenizer_X=None,
):
steps=0
e=flag
for item in train_loader:
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
GT=y_train_batch.cpu().detach()
GT= GT.unsqueeze(1)
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
print ('X_train_batch shape: ', X_train_batch.shape)
for iisample in range (len (cond_scales)):
if use_text_embedd:
result=model.sample (x= X_train_batch,stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample], device=device, skip_steps=skip_steps)
else:
result=model.sample (x= None, x_data_tokenized= X_train_batch,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample],device=device,skip_steps=skip_steps)
result=torch.round(result*ynormfac)
GT=torch.round (GT*ynormfac)
for samples in range (num_samples):
print ("sample ", samples, "out of ", num_samples)
plt.plot (result[samples,0,:].cpu().detach().numpy(),label= f'Predicted')
plt.plot (GT[samples,0,:],label= f'GT {0}')
plt.legend()
outname = prefix+ f"sample-{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
plt.savefig(outname, dpi=200)
plt.show ()
#reverse y sequence
to_rev=result[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
#reverse GT_y sequence
to_rev=GT[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
GT_y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
GT_y_data_reversed[iii]=GT_y_data_reversed[iii].upper().strip().replace(" ", "")
### reverse second structure input....
to_rev=torch.round (X_train_batch[:,:]*Xnormfac)
to_rev=to_rev.long().cpu().detach().numpy()
X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
print (f"For {X_train_batch[samples,:].cpu().detach().numpy()} or {X_data_reversed[samples]}, predicted sequence: ", y_data_reversed[samples])
print (f"Ground truth: {GT_y_data_reversed[samples]}")
if foldproteins:
xbc=X_train_batch[samples,:].cpu().detach().numpy()
out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})
tempname='temp'
pdb_file=foldandsavePDB (
sequence=y_data_reversed[samples],
filename_out=tempname,
num_cycle=16, flag=flag,
# +++++++++++++++++++
prefix=prefix
)
#out_nam=f'{prefix}{out_nam}.pdb'
out_nam=f'{prefix}{X_data_reversed[samples]}.pdb'
print (f'Original PDB: {pdb_file} OUT: {out_nam}')
shutil.copy (pdb_file, out_nam) #source, dest
pdb_file=out_nam
print (f"Properly named PDB file produced: {pdb_file}")
view=show_pdb(pdb_file=pdb_file, flag=flag, show_sidechains=show_sidechains, show_mainchains=show_mainchains, color=color)
view.show()
steps=steps+1
if steps>num_samples:
break
# ++
def sample_sequence_omegafold_pLM_ModelB (
model,
X=None, #this is the target conventionally when using text embd
flag=0,
cond_scales=1.,
foldproteins=False,
X_string=None,
x_data=None,
skip_steps=0,
inpaint_images = None,
inpaint_masks = None,
inpaint_resample_times = None,
init_images = None,
num_cycle=16,
# ++++++++++++++++++++++++
ynormfac=1,
train_unet_number=1,
tokenizer_X=None,
Xnormfac=1.,
max_length=1.,
prefix=None,
tokenizer_y=None,
# ++
CKeys=None,
sample_dir=None,
steps=None,
e=None,
IF_showfig=True, # effective only after foldproteins=True
# ++
pLM_Model=None, # pLM_Model,
pLM_Model_Name=None, # pLM_Model_Name,
image_channels=None, # image_channels,
pLM_alphabet=None, # esm_alphabet,
):
# steps=0
# e=flag
#num_samples = min (num_samples,y_train_batch.shape[0] )
if X!=None:
print (f"Producing {len(X)} samples...from text conditioning X...")
lenn_val=len(X)
if X_string!=None:
lenn_val=len(X_string)
print (f"Producing {len(X_string)} samples...from text conditioning X_String (from string)...")
if x_data!=None:
print (f"Producing {len(x_data)} samples...from image conditingig x_data ...")
lenn_val=len(x_data)
print (x_data)
print ('Device: ', device)
pdb_file_list=[]
fasta_file_list=[]
# + for debug
print('tot ', lenn_val)
for iisample in range (lenn_val):
print("Working on ", iisample)
X_cond=None # this is for text-conditioning
if X_string==None and X != None: #only do if X provided
X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
if X_string !=None:
XX = tokenizer_X.texts_to_sequences(X_string[iisample])
XX= sequence.pad_sequences(XX, maxlen=max_length, padding='post', truncating='post')
XX=np.array(XX)
X_cond=torch.from_numpy(XX).float()/Xnormfac
print ('Tokenized and processed: ', X_cond)
print ("X_cond=", X_cond)
# # --
# result=model.sample (
# x=X_cond,
# stop_at_unet_number=train_unet_number ,
# cond_scale=cond_scales ,
# x_data=x_data[iisample],
# # ++
# x_data_tokenized=
# skip_steps=skip_steps,
# inpaint_images = inpaint_images,
# inpaint_masks = inpaint_masks,
# inpaint_resample_times = inpaint_resample_times,
# init_images = init_images,device=device,
# # ++++++++++++++++++++++++++
# tokenizer_X=tokenizer_X,
# Xnormfac=Xnormfac,
# max_length=max_length,
# )
# ++
# use cond_image as the conditioning, via x_data_tokenized channel
# -----------------------------------------------------------------
# for below, two branches are all for cond_img, not for text_cond
if tokenizer_X!=None:
# for SecStr+ModelB
result_embedding=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales ,
x_data=x_data[iisample], # will pass through tokenizer_X in this sample(), channels will be matched with self.pred_dim
# ++
x_data_tokenized=None,
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,device=device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
)
else:
# for ForcPath+ModelB:
# for model.sample() here using x_data_tokenized channel
#
x_data_tokenized=torch.from_numpy(x_data[iisample]/Xnormfac)
x_data_tokenized=x_data_tokenized.to(torch.float)
# here, only one input list is read in
x_data_tokenized=x_data_tokenized.unsqueeze(0) # [batch=1, seq_len]
# leave channel expansion for the self.sample() to handle
# + for debug:
if CKeys['Debug_TrainerPack']==3:
print("x_data_tokenized dim: ", x_data_tokenized.shape)
print("x_data_tokenized dtype: ", x_data_tokenized.dtype)
print("test x_data_tokenized!=None: ", x_data_tokenized!=None)
result_embedding=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales ,
x_data=None,
# ++
x_data_tokenized=x_data_tokenized,
#
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,device=device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
)
# # --
# result=torch.round(result*ynormfac) # (batch=1, channel=1, seq_len)
# ++ for pLM
# full record
# result_embedding as image.dim: [batch, channels, seq_len]
# result_tokens.dim: [batch, seq_len]
if image_channels==33:
result_tokens,result_logits = convert_into_tokens_using_prob(
result_embedding,
pLM_Model_Name,
)
else:
result_tokens,result_logits = convert_into_tokens(
pLM_Model,
result_embedding,
pLM_Model_Name,
)
# +++++++++++++++++++++++++++++++++
result=result_tokens.unsqueeze(1) # dim: [batch, 1, seq_len]
# + for debug
print('result dim: ', result.shape)
# plot sequence token code: esm (33 tokens)
fig=plt.figure()
plt.plot (
result[0,0,:].cpu().detach().numpy(),
label= f'Predicted'
)
#plt.plot (GT[samples,0,:]*ynormfac,label= f'GT {0}')
plt.legend()
outname = sample_dir+ f"sampled_from_X_{iisample}_condscale-{str (cond_scales)}_{e}_{steps}.jpg"
#plt.title (f"Sample {samples}, cond scale={str (cond_scales[iisample])}")
if IF_showfig==1:
plt.show ()
else:
plt.savefig(outname, dpi=200)
plt.close()
#
# # --
# to_rev=result[:,0,:]
# to_rev=to_rev.long().cpu().detach().numpy()
# print (to_rev.shape)
# y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
#
# for iii in range (len(y_data_reversed)):
# y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
# # ++: for model A: no mask is provided from input
# # reverse the PREDICTED y into a foldable sequence
# # save this block for Model A
# to_rev=result[:,0,:] # token (batch,seq_len)
# y_data_reversed=decode_many_ems_token_rec_for_folding(
# to_rev,
# result_logits,
# pLM_alphabet,
# pLM_Model,
# )
# if CKeys['Debug_TrainerPack']==3:
# print("on foldable result: ", to_rev[0])
# print("on result_logits: ", result_logits[0])
# a = decode_one_ems_token_rec_for_folding(
# to_rev[0],
# result_logits[0],
# pLM_alphabet,
# pLM_Model,
# )
# print('One resu: ', a)
# print("on y_data_reversed: ", y_data_reversed[0])
# print("y_data_reversed.type", y_data_reversed.dtype)
#
# ++: for model B: using mask from the input
# extract the mask/seq_len from input if possible
if tokenizer_X!=None:
# for SecStr+ModelB
result_mask = read_mask_from_input(
tokenized_data=None,
mask_value=None,
seq_data=x_data[iisample],
max_seq_length=max_length,
)
else:
# for ForcPath+ModelB
result_mask = read_mask_from_input(
tokenized_data=x_data_tokenized, # None,
mask_value=0, # None,
seq_data=None, # x_data[iisample],
max_seq_length=None, # max_length,
)
to_rev=result[:,0,:] # token (batch,seq_len)
if CKeys['Debug_TrainerPack']==3:
print("on foldable result: ", to_rev[0])
print("on result_logits: ", result_logits[0])
print("on mask: ", result_mask[0])
a = decode_one_ems_token_rec_for_folding_with_mask(
to_rev[0],
result_logits[0],
pLM_alphabet,
pLM_Model,
result_mask[0],
)
print('One resu: ', a)
y_data_reversed=decode_many_ems_token_rec_for_folding_with_mask(
to_rev,
result_logits,
pLM_alphabet,
pLM_Model,
result_mask,
)
if CKeys['Debug_TrainerPack']==3:
print("on y_data_reversed[0]: ", y_data_reversed[0])
### reverse second structure input....
if X_cond != None:
X_cond=torch.round(X_cond*Xnormfac)
to_rev=X_cond[:,:]
to_rev=to_rev.long().cpu().detach().numpy()
print (to_rev.shape)
X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
if x_data !=None:
X_data_reversed=x_data #is already in sequence fromat..
# print (f"For {X} or {X_data_reversed[iisample]}, predicted sequence", y_data_reversed[iisample])
print (f"For {X} or {X_data_reversed[iisample]}, predicted sequence: ", y_data_reversed)
# + for debug
print("================================================")
print("foldproteins: ", foldproteins)
if not foldproteins:
pdb_file=None
else:
if X_cond != None:
xbc=X_cond[iisample,:].cpu().detach().numpy()
out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})+f'_{flag}_{steps}'
if x_data !=None:
#xbc=x_data[iisample]
# ----------------------------------
# this one can be too long for a name
out_nam=x_data[iisample]
# ++++++++++++++++++++++++++++++++++
#
out_nam=iisample
tempname='temp'
pdb_file, fasta_file=foldandsavePDB_pdb_fasta (
sequence=y_data_reversed[0],
filename_out=tempname,
num_cycle=num_cycle,
flag=flag,
# +++++++++++++++++++
# prefix=prefix,
prefix=sample_dir,
)
# out_nam_fasta=f'{prefix}{out_nam}_{flag}_{steps}.fasta'
# ------------------------------------------
# this one can be too long for a name
# out_nam_fasta=f'{sample_dir}{out_nam}_{flag}_{e}_{iisample}.fasta'
# ++++++++++++++++++++++++++++++++++++++++++
out_nam_fasta=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.fasta'
write_fasta (y_data_reversed[0], out_nam_fasta)
# out_nam=f'{prefix}{X_data_reversed[iisample]}_{flag}_{steps}.pdb'
# out_nam=f'{sample_dir}{X_data_reversed[iisample]}_{flag}_{steps}.pdb'
# -------------------------------------------
# this one can be too long for a name
# However, the input X is recorded in the code
# out_nam=f'{sample_dir}{X_data_reversed[iisample]}_{flag}_{iisample}.pdb'
# +++++++++++++++++++++++++++++++++++++++++++
out_nam=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.pdb'
out_nam_fasta=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.fasta'
# print('Debug 1: out: ', out_nam)
# print('Debug 2: in: ', pdb_file)
shutil.copy (pdb_file, out_nam) #source, dest
shutil.copy (fasta_file, out_nam_fasta)
# cmd_line = 'cp ' + pdb_file + ' ' + out_nam
# print(cmd_line)
# os.popen(cmd_line)
# print('Debug 3')
# clean the slade to avoid mistakenly using the previous fasta file
os.remove (pdb_file)
os.remove (fasta_file)
pdb_file=out_nam
fasta_file=out_nam_fasta
pdb_file_list.append(pdb_file)
fasta_file_list.append(fasta_file)
print (f"Properly named PDB file produced: {pdb_file}")
if IF_showfig==1:
#flag=1000
view=show_pdb(
pdb_file=pdb_file,
flag=flag,
show_sidechains=show_sidechains,
show_mainchains=show_mainchains,
color=color
)
view.show()
# steps=steps+1
return pdb_file_list, fasta_file_list
#
def sample_sequence_omegafold_ModelB (
model,
X=None, #this is the target conventionally when using text embd
flag=0,
cond_scales=1.,
foldproteins=False,
X_string=None,
x_data=None,
skip_steps=0,
inpaint_images = None,
inpaint_masks = None,
inpaint_resample_times = None,
init_images = None,
num_cycle=16,
# ++++++++++++++++++++++++
ynormfac=1,
train_unet_number=1,
tokenizer_X=None,
Xnormfac=1.,
max_length=1.,
prefix=None,
tokenizer_y=None,
# ++
CKeys=None,
sample_dir=None,
steps=None,
e=None,
IF_showfig=True, # effective only after foldproteins=True
):
# steps=0
# e=flag
#num_samples = min (num_samples,y_train_batch.shape[0] )
if X!=None:
print (f"Producing {len(X)} samples...from text conditioning X...")
lenn_val=len(X)
if X_string!=None:
lenn_val=len(X_string)
print (f"Producing {len(X_string)} samples...from text conditioning X_String (from string)...")
if x_data!=None:
print (f"Producing {len(x_data)} samples...from image conditingig x_data ...")
lenn_val=len(x_data)
print (x_data)
print ('Device: ', device)
# + for debug
print('tot ', lenn_val)
for iisample in range (lenn_val):
print("Working on ", iisample)
X_cond=None
if X_string==None and X != None: #only do if X provided
X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
if X_string !=None:
XX = tokenizer_X.texts_to_sequences(X_string[iisample])
XX= sequence.pad_sequences(XX, maxlen=max_length, padding='post', truncating='post')
XX=np.array(XX)
X_cond=torch.from_numpy(XX).float()/Xnormfac
print ('Tokenized and processed: ', X_cond)
print ("X_cond=", X_cond)
# # --
# result=model.sample (
# x=X_cond,
# stop_at_unet_number=train_unet_number ,
# cond_scale=cond_scales ,
# x_data=x_data[iisample],
# # ++
# x_data_tokenized=
# skip_steps=skip_steps,
# inpaint_images = inpaint_images,
# inpaint_masks = inpaint_masks,
# inpaint_resample_times = inpaint_resample_times,
# init_images = init_images,device=device,
# # ++++++++++++++++++++++++++
# tokenizer_X=tokenizer_X,
# Xnormfac=Xnormfac,
# max_length=max_length,
# )
# ++
if tokenizer_X!=None:
result=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales ,
x_data=x_data[iisample],
# ++
x_data_tokenized=None,
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,device=device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
)
else:
x_data_tokenized=torch.from_numpy(x_data[iisample]/Xnormfac)
x_data_tokenized=x_data_tokenized.to(torch.float)
# + for debug:
if CKeys['Debug_TrainerPack']==1:
print("x_data_tokenized dim: ", x_data_tokenized.shape)
print("x_data_tokenized dtype: ", x_data_tokenized.dtype)
print("test: ", x_data_tokenized!=None)
result=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales ,
x_data=None,
# ++
x_data_tokenized=x_data_tokenized,
#
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,device=device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
)
result=torch.round(result*ynormfac)
# + for debug
print('result dim: ', result.shape)
fig=plt.figure()
plt.plot (
result[0,0,:].cpu().detach().numpy(),
label= f'Predicted'
)
#plt.plot (GT[samples,0,:]*ynormfac,label= f'GT {0}')
plt.legend()
outname = sample_dir+ f"sampled_from_X_{iisample}_condscale-{str (cond_scales)}_{e}_{steps}.jpg"
#plt.title (f"Sample {samples}, cond scale={str (cond_scales[iisample])}")
if IF_showfig==1:
plt.show ()
else:
plt.savefig(outname, dpi=200)
plt.close()
to_rev=result[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
print (to_rev.shape)
y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
### reverse second structure input....
if X_cond != None:
X_cond=torch.round(X_cond*Xnormfac)
to_rev=X_cond[:,:]
to_rev=to_rev.long().cpu().detach().numpy()
print (to_rev.shape)
X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
if x_data !=None:
X_data_reversed=x_data #is already in sequence fromat..
# print (f"For {X} or {X_data_reversed[iisample]}, predicted sequence", y_data_reversed[iisample])
print (f"For {X} or {X_data_reversed[iisample]}, predicted sequence: ", y_data_reversed)
# + for debug
print("================================================")
print("foldproteins: ", foldproteins)
if not foldproteins:
pdb_file=None
else:
if X_cond != None:
xbc=X_cond[iisample,:].cpu().detach().numpy()
out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})+f'_{flag}_{steps}'
if x_data !=None:
#xbc=x_data[iisample]
# ----------------------------------
# this one can be too long for a name
out_nam=x_data[iisample]
# ++++++++++++++++++++++++++++++++++
#
out_nam=iisample
tempname='temp'
pdb_file=foldandsavePDB (
sequence=y_data_reversed[0],
filename_out=tempname,
num_cycle=num_cycle,
flag=flag,
# +++++++++++++++++++
# prefix=prefix,
prefix=sample_dir,
)
# out_nam_fasta=f'{prefix}{out_nam}_{flag}_{steps}.fasta'
# ------------------------------------------
# this one can be too long for a name
# out_nam_fasta=f'{sample_dir}{out_nam}_{flag}_{e}_{iisample}.fasta'
# ++++++++++++++++++++++++++++++++++++++++++
out_nam_fasta=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.fasta'
write_fasta (y_data_reversed[0], out_nam_fasta)
# out_nam=f'{prefix}{X_data_reversed[iisample]}_{flag}_{steps}.pdb'
# out_nam=f'{sample_dir}{X_data_reversed[iisample]}_{flag}_{steps}.pdb'
# -------------------------------------------
# this one can be too long for a name
# However, the input X is recorded in the code
# out_nam=f'{sample_dir}{X_data_reversed[iisample]}_{flag}_{iisample}.pdb'
# +++++++++++++++++++++++++++++++++++++++++++
out_nam=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.pdb'
# print('Debug 1: out: ', out_nam)
# print('Debug 2: in: ', pdb_file)
shutil.copy (pdb_file, out_nam) #source, dest
# cmd_line = 'cp ' + pdb_file + ' ' + out_nam
# print(cmd_line)
# os.popen(cmd_line)
# print('Debug 3')
pdb_file=out_nam
print (f"Properly named PDB file produced: {pdb_file}")
if IF_showfig==1:
#flag=1000
view=show_pdb(
pdb_file=pdb_file,
flag=flag,
show_sidechains=show_sidechains,
show_mainchains=show_mainchains,
color=color
)
view.show()
# steps=steps+1
return pdb_file
# ++ for de novo input of ForcPath
# ++
def sample_sequence_omegafold_pLM_ModelB_For_ForcPath (
model,
X=None, #this is the target conventionally when using text embd
flag=0,
cond_scales=[1.], # 1.,
foldproteins=False,
X_string=None,
x_data=None,
skip_steps=0,
inpaint_images = None,
inpaint_masks = None,
inpaint_resample_times = None,
init_images = None,
num_cycle=16,
# ++++++++++++++++++++++++
ynormfac=1,
train_unet_number=1,
tokenizer_X=None,
Xnormfac=1.,
max_length=1.,
prefix=None,
tokenizer_y=None,
# ++
CKeys=None,
sample_dir=None,
steps=None,
e=None,
IF_showfig=True, # effective only after foldproteins=True
# ++
pLM_Model=None, # pLM_Model,
pLM_Model_Name=None, # pLM_Model_Name,
image_channels=None, # image_channels,
pLM_alphabet=None, # esm_alphabet,
):
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# prepare input in different channels
#
X_cond=None # this is for text-conditioning
if X_string==None and X != None: #only do if X provided
print (f"Producing {len(X)} samples...from text conditioning X...")
lenn_val=len(X)
# shape of X: [[..],[..]]: double bracket
X_cond=torch.Tensor(X).to(device)
# --
# X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
if X_string !=None:
print (f"Producing {len(X_string)} samples...from text conditioning X_String (from string)...")
lenn_val=len(X_string)
# --
XX = tokenizer_X.texts_to_sequences(X_string[iisample])
# ++
XX = tokenizer_X.texts_to_sequences(X_string)
XX= sequence.pad_sequences(XX, maxlen=max_length, padding='post', truncating='post')
XX=np.array(XX)
X_cond=torch.from_numpy(XX).float()/Xnormfac
print ('Tokenized and processed: ', X_cond)
if x_data!=None:
print (f"Producing {len(x_data)} samples...from image conditingig x_data ...")
lenn_val=len(x_data)
if tokenizer_X==None: # for ForcPath,
# need to do Padding and Normalization
# and then put into tokenized data channel
x_data_tokenized=[]
for ii in range(lenn_val):
x_data_one_line=pad_a_np_arr(x_data[ii], 0.0, max_length)
x_data_tokenized.append(x_data_one_line)
x_data_tokenized=np.array(x_data_tokenized)
x_data_tokenized=torch.from_numpy(x_data_tokenized/Xnormfac)
else:
# leave for SecStr case: TBA
pass
# print (x_data)
# ++ for result_mask based on input: x_data or x_data_tokenized
# ++: for model B: using mask from the input
# extract the mask/seq_len from input if possible
if tokenizer_X!=None:
# for SecStr+ModelB
result_mask = read_mask_from_input(
tokenized_data=None,
mask_value=None,
seq_data=x_data, # x_data[iisample],
max_seq_length=max_length,
)
else:
# for ForcPath+ModelB
result_mask = read_mask_from_input(
tokenized_data=x_data_tokenized, # None,
mask_value=0, # None,
seq_data=None, # x_data[iisample],
max_seq_length=None, # max_length,
)
print ("Input contents:")
print ("cond_img condition: x_data=\n", x_data)
print ("Text condition: X_cond=\n", X_cond)
# store the results
pdb_file_list=[]
fasta_file_list=[]
# loop over cond_scales
for idx_cond, this_cond_scale in enumerate(cond_scales):
print(f"Working on cond_scale {str(this_cond_scale)}")
# do sampling
# -----------------------------------------------------------------
# for below, two branches are all for cond_img, not for text_cond
if tokenizer_X!=None:
# for SecStr+ModelB, not test here
result_embedding=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=this_cond_scale, # cond_scales ,
x_data=x_data, # x_data[iisample], # will pass through tokenizer_X in this sample(), channels will be matched with self.pred_dim
# ++
x_data_tokenized=None,
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,device=device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
)
else:
# for ForcPath+ModelB:
# for model.sample() here using x_data_tokenized channel
x_data_tokenized=x_data_tokenized.to(torch.float) # shape [batch, max_seq_len]
# leave channel expansion for the self.sample() to handle
# + for debug:
if CKeys['Debug_TrainerPack']==3:
print("x_data_tokenized dim: ", x_data_tokenized.shape)
print("x_data_tokenized dtype: ", x_data_tokenized.dtype)
print("test x_data_tokenized!=None: ", x_data_tokenized!=None)
result_embedding=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=this_cond_scale, # cond_scales ,
x_data=None,
# ++
x_data_tokenized=x_data_tokenized,
#
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,device=device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
)
# handle the results: from embedding into AA
# ++ for pLM
if image_channels==33:
# pass
result_tokens,result_logits = convert_into_tokens_using_prob(
result_embedding,
pLM_Model_Name,
)
else:
# full record
# result_embedding as image.dim: [batch, channels, seq_len]
# result_tokens.dim: [batch, seq_len]
result_tokens,result_logits = convert_into_tokens(
pLM_Model,
result_embedding,
pLM_Model_Name,
)
# +++++++++++++++++++++++++++++++++
result=result_tokens.unsqueeze(1) # dim: [batch, 1, seq_len]
# + for debug
print('result dim: ', result.shape)
# plot sequence token code: esm (33 tokens), for one batch
fig=plt.figure()
for ii in range(lenn_val):
plt.plot (
result[ii,0,:].cpu().detach().numpy(),
label= f'Predicted for Input#{str(ii)}'
)
#plt.plot (GT[samples,0,:]*ynormfac,label= f'GT {0}')
plt.legend()
outname = sample_dir+ f"DenovoInputXs_CondScale_No{str(idx_cond)}_Val_{str(this_cond_scale)}_{e}_{steps}.jpg"
#plt.title (f"Sample {samples}, cond scale={str (cond_scales[iisample])}")
if IF_showfig==1:
plt.show ()
else:
plt.savefig(outname, dpi=200)
plt.close()
# translate result into AA
to_rev=result[:,0,:] # token (batch,seq_len)
if CKeys['Debug_TrainerPack']==3:
print("on foldable result: ", to_rev[0])
print("on result_logits: ", result_logits[0])
print("on mask: ", result_mask[0])
a = decode_one_ems_token_rec_for_folding_with_mask(
to_rev[0],
result_logits[0],
pLM_alphabet,
pLM_Model,
result_mask[0],
)
print('One resu: ', a)
y_data_reversed=decode_many_ems_token_rec_for_folding_with_mask(
to_rev,
result_logits,
pLM_alphabet,
pLM_Model,
result_mask,
)
if CKeys['Debug_TrainerPack']==3:
print("on y_data_reversed[0]: ", y_data_reversed[0])
### reverse second structure input....
if X_cond != None:
X_cond=torch.round(X_cond*Xnormfac)
to_rev=X_cond[:,:]
to_rev=to_rev.long().cpu().detach().numpy()
print (to_rev.shape)
X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
if x_data !=None:
# work for second structure input....
# work for ForcPath input...
X_data_reversed=x_data #is already in sequence fromat..
# sections for each one result
for iisample in range(lenn_val):
print (f"For {X} or {X_data_reversed[iisample]}, predicted sequence: ", y_data_reversed[iisample])
out_nam_fasta=f'{sample_dir}DN_{iisample}_CondS_No_{idx_cond}_Val_{this_cond_scale}_epo_{e}_step_{steps}.fasta'
write_fasta (y_data_reversed[iisample], out_nam_fasta)
fasta_file_list.append(out_nam_fasta)
# + for debug
print("================================================")
print("foldproteins: ", foldproteins)
if not foldproteins:
pdb_file=None
else:
if X_cond != None:
# not maintained
xbc=X_cond[iisample,:].cpu().detach().numpy()
out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.4f" % xbc})+f'_{flag}_{steps}'
if x_data !=None:
pass
# #xbc=x_data[iisample]
# # ----------------------------------
# # this one can be too long for a name
# out_nam=x_data[iisample]
# # ++++++++++++++++++++++++++++++++++
# #
# out_nam=iisample
tempname='temp'
pdb_file, fasta_file=foldandsavePDB_pdb_fasta (
sequence=y_data_reversed[iisample],
filename_out=tempname,
num_cycle=num_cycle,
flag=flag,
# +++++++++++++++++++
# prefix=prefix,
prefix=sample_dir,
)
# out_nam=f'{prefix}{X_data_reversed[iisample]}_{flag}_{steps}.pdb'
# out_nam=f'{sample_dir}{X_data_reversed[iisample]}_{flag}_{steps}.pdb'
# -------------------------------------------
# this one can be too long for a name
# However, the input X is recorded in the code
# out_nam=f'{sample_dir}{X_data_reversed[iisample]}_{flag}_{iisample}.pdb'
# +++++++++++++++++++++++++++++++++++++++++++
out_nam=f'{sample_dir}DN_{iisample}_CondS_No_{idx_cond}_Val_{this_cond_scale}_epo_{e}_step_{steps}.pdb'
# out_nam_fasta=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.fasta'
# print('Debug 1: out: ', out_nam)
# print('Debug 2: in: ', pdb_file)
shutil.copy (pdb_file, out_nam) #source, dest
# shutil.copy (fasta_file, out_nam_fasta)
# cmd_line = 'cp ' + pdb_file + ' ' + out_nam
# print(cmd_line)
# os.popen(cmd_line)
# print('Debug 3')
# clean the slade to avoid mistakenly using the previous fasta file
os.remove (pdb_file)
os.remove (fasta_file)
pdb_file=out_nam
# fasta_file=out_nam_fasta
pdb_file_list.append(pdb_file)
# ++ write the input condtion as a reference: for ForcPath
out_nam_inX=f'{sample_dir}DN_{iisample}_CondS_No_{idx_cond}_Val_{this_cond_scale}_epo_{e}_step_{steps}_input.txt'
if torch.is_tensor(X_data_reversed[iisample]):
# for safety, not used usually
xbc=X_data_reversed[iisample].cpu().detach().numpy()
else:
xbc=X_data_reversed[iisample]
if tokenizer_X==None:
# for ForcPath case
out_inX=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.4f" % xbc})
else:
# for SecStr case
out_inX=xbc
with open(out_nam_inX, "w") as inX_file:
inX_file.write(out_inX)
print (f"Properly named PDB file produced: {pdb_file}")
if IF_showfig==1:
#flag=1000
view=show_pdb(
pdb_file=pdb_file,
flag=flag,
show_sidechains=show_sidechains,
show_mainchains=show_mainchains,
color=color
)
view.show()
return pdb_file_list, fasta_file_list
# ++
def sample_sequence_pLM_ModelB_For_ForcPath_Predictor (
model,
X=None, #this is the target conventionally when using text embd
flag=0,
cond_scales=[1.], # 1.,
foldproteins=False,
X_string=None,
x_data=None,
skip_steps=0,
inpaint_images = None,
inpaint_masks = None,
inpaint_resample_times = None,
init_images = None,
num_cycle=16,
# ++++++++++++++++++++++++
ynormfac=1,
train_unet_number=1,
tokenizer_X=None,
Xnormfac=1.,
max_length=1.,
prefix=None,
tokenizer_y=None,
# ++
CKeys=None,
sample_dir=None,
steps=None,
e=None,
IF_showfig=True, # effective only after foldproteins=True
# ++
pLM_Model=None, # pLM_Model,
pLM_Model_Name=None, # pLM_Model_Name,
image_channels=None, # image_channels,
pLM_alphabet=None, # esm_alphabet,
# ++
esm_layer=None,
):
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# input: a list of AA sequence in string format
# output: ForcPath prediction
#
# 1. decide input channel: text_cond or img_cond
X_cond=None # this is for text-conditioning
if X_string==None and X != None: #only do if X provided
print (f"Producing {len(X)} samples...from text conditioning X...")
lenn_val=len(X)
# shape of X: [[..],[..]]: double bracket
X_cond=torch.Tensor(X).to(device)
# --
# X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
if X_string !=None:
print (f"Producing {len(X_string)} samples...from text conditioning X_String (from string)...")
lenn_val=len(X_string)
# --
XX = tokenizer_X.texts_to_sequences(X_string[iisample])
# ++
XX = tokenizer_X.texts_to_sequences(X_string)
XX= sequence.pad_sequences(XX, maxlen=max_length, padding='post', truncating='post')
XX=np.array(XX)
X_cond=torch.from_numpy(XX).float()/Xnormfac
print ('Tokenized and processed: ', X_cond)
if x_data!=None: # this is for img_conditioning channel
# Will use this channel
print (f"Producing {len(x_data)} samples...from image conditingig x_data ...")
lenn_val=len(x_data)
seq_len_list=[]
for this_AA in x_data:
seq_len_list.append(len(this_AA))
print ("Input contents:")
print ("cond_img condition: x_data=\n", x_data)
print ("Text condition: X_cond=\n", X_cond)
# 2. perform sampling
# loop over cond_scales
resu_prediction={}
for idx_cond, this_cond_scale in enumerate(cond_scales):
print(f"Working on cond_scale {str(this_cond_scale)}")
# leave the translation from seq to tokenized
# in the model.sample function using x_data channel
# Need to pass on the esm model part or tokenizer_X
#
result_embedding=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=this_cond_scale, # cond_scales ,
x_data=x_data, # x_data[iisample], # will pass through tokenizer_X in this sample(), channels will be matched with self.pred_dim
# ++
x_data_tokenized=None,
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,device=device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
# ++ for esm
pLM_Model=pLM_Model,
pLM_alphabet=pLM_alphabet,
esm_layer=esm_layer,
pLM_Model_Name=pLM_Model_Name,
# image_channels=image_channels,
)
# convert into prediction
# 3. translate prediction into something meaningful
# consider channel average and masking
# average across channels
result_embedding=torch.mean(result_embedding, 1) # (batch, seq_len)
# read mask from input: X_train_batch_picked (batch, seq_len)
# result_mask looks like, 0,1,1,...,1,0,0
# will fill 0th component be zero
result_mask = read_mask_from_input(
tokenized_data=None, # X_train_batch[:num_samples],
mask_value=0.0,
seq_data=x_data, # None,
max_seq_length=max_length, # None,
)
# apply mask to result: keep true and zero all false
# this also make sure 0th components are zero ACCIDENTLY
result = result_embedding.cpu()*result_mask # (batch, seq_len)
# result = result.cpu()
y_data_reversed = result*ynormfac
# 4. translate the results into a list
prediction_list = []
for ii in range(len(x_data)):
prediction_list.append(
y_data_reversed[ii, :seq_len_list[ii]+1]
)
if CKeys['Debug_TrainerPack']==3:
print("check prediction dim:")
print(f"model output: ", y_data_reversed[0])
print(f"prediction output: ", prediction_list[0])
#
# store the results
resu_prediction[str(this_cond_scale)]=prediction_list
return resu_prediction,seq_len_list
# # --------------------------------------------------------------------------
# # prepare input in different channels
# #
# X_cond=None # this is for text-conditioning
# if X_string==None and X != None: #only do if X provided
# print (f"Producing {len(X)} samples...from text conditioning X...")
# lenn_val=len(X)
# # shape of X: [[..],[..]]: double bracket
# X_cond=torch.Tensor(X).to(device)
# # --
# # X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
# if X_string !=None:
# print (f"Producing {len(X_string)} samples...from text conditioning X_String (from string)...")
# lenn_val=len(X_string)
# # --
# XX = tokenizer_X.texts_to_sequences(X_string[iisample])
# # ++
# XX = tokenizer_X.texts_to_sequences(X_string)
# XX= sequence.pad_sequences(XX, maxlen=max_length, padding='post', truncating='post')
# XX=np.array(XX)
# X_cond=torch.from_numpy(XX).float()/Xnormfac
# print ('Tokenized and processed: ', X_cond)
# if x_data!=None:
# print (f"Producing {len(x_data)} samples...from image conditingig x_data ...")
# lenn_val=len(x_data)
# if tokenizer_X==None: # for ForcPath,
# # need to do Padding and Normalization
# # and then put into tokenized data channel
# x_data_tokenized=[]
# for ii in range(lenn_val):
# x_data_one_line=pad_a_np_arr(x_data[ii], 0.0, max_length)
# x_data_tokenized.append(x_data_one_line)
# x_data_tokenized=np.array(x_data_tokenized)
# x_data_tokenized=torch.from_numpy(x_data_tokenized/Xnormfac)
# else:
# # leave for SecStr case: TBA
# pass
# # print (x_data)
# # ++ for result_mask based on input: x_data or x_data_tokenized
# # ++: for model B: using mask from the input
# # extract the mask/seq_len from input if possible
# if tokenizer_X!=None:
# # for SecStr+ModelB
# result_mask = read_mask_from_input(
# tokenized_data=None,
# mask_value=None,
# seq_data=x_data, # x_data[iisample],
# max_seq_length=max_length,
# )
# else:
# # for ForcPath+ModelB
# result_mask = read_mask_from_input(
# tokenized_data=x_data_tokenized, # None,
# mask_value=0, # None,
# seq_data=None, # x_data[iisample],
# max_seq_length=None, # max_length,
# )
# print ("Input contents:")
# print ("cond_img condition: x_data=\n", x_data)
# print ("Text condition: X_cond=\n", X_cond)
# # store the results
# pdb_file_list=[]
# fasta_file_list=[]
# # loop over cond_scales
# for idx_cond, this_cond_scale in enumerate(cond_scales):
# print(f"Working on cond_scale {str(this_cond_scale)}")
# # do sampling
# # -----------------------------------------------------------------
# # for below, two branches are all for cond_img, not for text_cond
# if tokenizer_X!=None:
# # for SecStr+ModelB, not test here
# result_embedding=model.sample (
# x=X_cond,
# stop_at_unet_number=train_unet_number ,
# cond_scale=this_cond_scale, # cond_scales ,
# x_data=x_data, # x_data[iisample], # will pass through tokenizer_X in this sample(), channels will be matched with self.pred_dim
# # ++
# x_data_tokenized=None,
# skip_steps=skip_steps,
# inpaint_images = inpaint_images,
# inpaint_masks = inpaint_masks,
# inpaint_resample_times = inpaint_resample_times,
# init_images = init_images,device=device,
# # ++++++++++++++++++++++++++
# tokenizer_X=tokenizer_X,
# Xnormfac=Xnormfac,
# max_length=max_length,
# )
# else:
# # for ForcPath+ModelB:
# # for model.sample() here using x_data_tokenized channel
# x_data_tokenized=x_data_tokenized.to(torch.float) # shape [batch, max_seq_len]
# # leave channel expansion for the self.sample() to handle
# # + for debug:
# if CKeys['Debug_TrainerPack']==3:
# print("x_data_tokenized dim: ", x_data_tokenized.shape)
# print("x_data_tokenized dtype: ", x_data_tokenized.dtype)
# print("test x_data_tokenized!=None: ", x_data_tokenized!=None)
# result_embedding=model.sample (
# x=X_cond,
# stop_at_unet_number=train_unet_number ,
# cond_scale=this_cond_scale, # cond_scales ,
# x_data=None,
# # ++
# x_data_tokenized=x_data_tokenized,
# #
# skip_steps=skip_steps,
# inpaint_images = inpaint_images,
# inpaint_masks = inpaint_masks,
# inpaint_resample_times = inpaint_resample_times,
# init_images = init_images,device=device,
# # ++++++++++++++++++++++++++
# tokenizer_X=tokenizer_X,
# Xnormfac=Xnormfac,
# max_length=max_length,
# )
# # handle the results: from embedding into AA
# # ++ for pLM
# # full record
# # result_embedding as image.dim: [batch, channels, seq_len]
# # result_tokens.dim: [batch, seq_len]
# result_tokens,result_logits = convert_into_tokens(
# pLM_Model,
# result_embedding,
# pLM_Model_Name,
# )
# # +++++++++++++++++++++++++++++++++
# result=result_tokens.unsqueeze(1) # dim: [batch, 1, seq_len]
# # + for debug
# print('result dim: ', result.shape)
# # plot sequence token code: esm (33 tokens), for one batch
# fig=plt.figure()
# for ii in range(lenn_val):
# plt.plot (
# result[ii,0,:].cpu().detach().numpy(),
# label= f'Predicted for Input#{str(ii)}'
# )
# #plt.plot (GT[samples,0,:]*ynormfac,label= f'GT {0}')
# plt.legend()
# outname = sample_dir+ f"DenovoInputXs_CondScale_No{str(idx_cond)}_Val_{str(this_cond_scale)}_{e}_{steps}.jpg"
# #plt.title (f"Sample {samples}, cond scale={str (cond_scales[iisample])}")
# if IF_showfig==1:
# plt.show ()
# else:
# plt.savefig(outname, dpi=200)
# plt.close()
# # translate result into AA
# to_rev=result[:,0,:] # token (batch,seq_len)
# if CKeys['Debug_TrainerPack']==3:
# print("on foldable result: ", to_rev[0])
# print("on result_logits: ", result_logits[0])
# print("on mask: ", result_mask[0])
# a = decode_one_ems_token_rec_for_folding_with_mask(
# to_rev[0],
# result_logits[0],
# pLM_alphabet,
# pLM_Model,
# result_mask[0],
# )
# print('One resu: ', a)
# y_data_reversed=decode_many_ems_token_rec_for_folding_with_mask(
# to_rev,
# result_logits,
# pLM_alphabet,
# pLM_Model,
# result_mask,
# )
# if CKeys['Debug_TrainerPack']==3:
# print("on y_data_reversed[0]: ", y_data_reversed[0])
# ### reverse second structure input....
# if X_cond != None:
# X_cond=torch.round(X_cond*Xnormfac)
# to_rev=X_cond[:,:]
# to_rev=to_rev.long().cpu().detach().numpy()
# print (to_rev.shape)
# X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
# for iii in range (len(y_data_reversed)):
# X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
# if x_data !=None:
# # work for second structure input....
# # work for ForcPath input...
# X_data_reversed=x_data #is already in sequence fromat..
# # sections for each one result
# for iisample in range(lenn_val):
# print (f"For {X} or {X_data_reversed[iisample]}, predicted sequence: ", y_data_reversed[iisample])
# out_nam_fasta=f'{sample_dir}DN_{iisample}_CondS_No_{idx_cond}_Val_{this_cond_scale}_epo_{e}_step_{steps}.fasta'
# write_fasta (y_data_reversed[iisample], out_nam_fasta)
# fasta_file_list.append(out_nam_fasta)
# # + for debug
# print("================================================")
# print("foldproteins: ", foldproteins)
# if not foldproteins:
# pdb_file=None
# else:
# if X_cond != None:
# # not maintained
# xbc=X_cond[iisample,:].cpu().detach().numpy()
# out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.4f" % xbc})+f'_{flag}_{steps}'
# if x_data !=None:
# pass
# # #xbc=x_data[iisample]
# # # ----------------------------------
# # # this one can be too long for a name
# # out_nam=x_data[iisample]
# # # ++++++++++++++++++++++++++++++++++
# # #
# # out_nam=iisample
# tempname='temp'
# pdb_file, fasta_file=foldandsavePDB_pdb_fasta (
# sequence=y_data_reversed[iisample],
# filename_out=tempname,
# num_cycle=num_cycle,
# flag=flag,
# # +++++++++++++++++++
# # prefix=prefix,
# prefix=sample_dir,
# )
# # out_nam=f'{prefix}{X_data_reversed[iisample]}_{flag}_{steps}.pdb'
# # out_nam=f'{sample_dir}{X_data_reversed[iisample]}_{flag}_{steps}.pdb'
# # -------------------------------------------
# # this one can be too long for a name
# # However, the input X is recorded in the code
# # out_nam=f'{sample_dir}{X_data_reversed[iisample]}_{flag}_{iisample}.pdb'
# # +++++++++++++++++++++++++++++++++++++++++++
# out_nam=f'{sample_dir}DN_{iisample}_CondS_No_{idx_cond}_Val_{this_cond_scale}_epo_{e}_step_{steps}.pdb'
# # out_nam_fasta=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.fasta'
# # print('Debug 1: out: ', out_nam)
# # print('Debug 2: in: ', pdb_file)
# shutil.copy (pdb_file, out_nam) #source, dest
# # shutil.copy (fasta_file, out_nam_fasta)
# # cmd_line = 'cp ' + pdb_file + ' ' + out_nam
# # print(cmd_line)
# # os.popen(cmd_line)
# # print('Debug 3')
# # clean the slade to avoid mistakenly using the previous fasta file
# os.remove (pdb_file)
# os.remove (fasta_file)
# pdb_file=out_nam
# # fasta_file=out_nam_fasta
# pdb_file_list.append(pdb_file)
# # ++ write the input condtion as a reference: for ForcPath
# out_nam_inX=f'{sample_dir}DN_{iisample}_CondS_No_{idx_cond}_Val_{this_cond_scale}_epo_{e}_step_{steps}_input.txt'
# if torch.is_tensor(X_data_reversed[iisample]):
# # for safety, not used usually
# xbc=X_data_reversed[iisample].cpu().detach().numpy()
# else:
# xbc=X_data_reversed[iisample]
# if tokenizer_X==None:
# # for ForcPath case
# out_inX=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.4f" % xbc})
# else:
# # for SecStr case
# out_inX=xbc
# with open(out_nam_inX, "w") as inX_file:
# inX_file.write(out_inX)
# print (f"Properly named PDB file produced: {pdb_file}")
# if IF_showfig==1:
# #flag=1000
# view=show_pdb(
# pdb_file=pdb_file,
# flag=flag,
# show_sidechains=show_sidechains,
# show_mainchains=show_mainchains,
# color=color
# )
# view.show()
# return pdb_file_list, fasta_file_list
# ++
# for ProteinDesigner
def sample_loop_omegafold_pLM_ModelB (
model,
train_loader,
cond_scales=[7.5], #list of cond scales - each sampled...
num_samples=2, #how many samples produced every time tested.....
timesteps=100, # not used
flag=0,
foldproteins=False,
use_text_embedd=True,
skip_steps=0,
# +++++++++++++++++++
train_unet_number=1,
ynormfac=1,
prefix=None,
tokenizer_y=None,
Xnormfac=1,
tokenizer_X=None,
# ++
CKeys=None,
sample_dir=None,
steps=None,
e=None,
IF_showfig=True, # effective only after foldproteins=True
# ++
pLM_Model=None,
pLM_Model_Name=None,
image_channels=None,
pLM_alphabet=None,
):
# =====================================================
# sample # = num_samples*(# of mini-batches)
# =====================================================
# steps=0
# e=flag
# for item in train_loader:
for idx, item in enumerate(train_loader):
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
# --
# # ++ for pLM case:
# if pLM_Model_Name=='None':
# # just use the encoded sequence
# # y_train_batch_in = y_train_batch.unsqueeze(1)
# X_train_batch_in = X_train_batch.unsqueeze(1)
# # pass
# elif pLM_Model_Name=='esm2_t33_650M_UR50D':
# # with torch.no_grad():
# # results = pLM_Model(
# # y_train_batch,
# # repr_layers=[33],
# # return_contacts=False,
# # )
# # y_train_batch_in = results["representations"][33]
# # y_train_batch_in = rearrange(
# # y_train_batch_in,
# # 'b l c -> b c l'
# # )
# X_train_batch_in = X_train_batch.unsqueeze(1).repeat(1,image_channels,1)
GT=y_train_batch.cpu().detach()
GT= GT.unsqueeze(1)
if num_samples>y_train_batch.shape[0]:
print("Warning: sampling # > len(mini_batch)")
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
X_train_batch_picked = X_train_batch[:num_samples,:] # X_train_batch_in[:num_samples ] #
print ('After pLM, (TEST) X_batch shape: ', X_train_batch_picked.shape)
for iisample in range (len (cond_scales)):
if use_text_embedd:
result_embedding=model.sample (
# x= X_train_batch,
x= X_train_batch_picked,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample],
device=device,
skip_steps=skip_steps
)
else:
result_embedding=model.sample (
x= None,
# x_data_tokenized= X_train_batch,
x_data_tokenized= X_train_batch_picked, # dim=(batch, seq_len), will extend channels inside .sample(),
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample],
device=device,
skip_steps=skip_steps
)
# ++ for pLM:
if image_channels==33:
result_tokens,result_logits = convert_into_tokens_using_prob(
result_embedding,
pLM_Model_Name,
)
else:
# full record
# result_embedding as image.dim: [batch, channels, seq_len]
# result_tokens.dim: [batch, seq_len]
result_tokens,result_logits = convert_into_tokens(
pLM_Model,
result_embedding,
pLM_Model_Name,
)
# # ---------------------------------
# result=torch.round(result*ynormfac)
# GT=torch.round (GT*ynormfac)
# +++++++++++++++++++++++++++++++++
result=result_tokens.unsqueeze(1) # dim: [batch, 1, seq_len]
# +
# # -------------------------------------------
# #reverse y sequence
# to_rev=result[:,0,:]
# to_rev=to_rev.long().cpu().detach().numpy()
#
# y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
#
# for iii in range (len(y_data_reversed)):
# y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
# ++++++++++++++++++++++++++++++++++++++++++++
# extract the mask/seq_len from input if possible
# here, from dataloader, we only use tokenized_data for mask generation
result_mask = read_mask_from_input(
tokenized_data=X_train_batch[:num_samples],
mask_value=0,
seq_data=None,
max_seq_length=None,
)
to_rev=result[:,0,:] # token (batch,seq_len)
if CKeys['Debug_TrainerPack']==3:
print("on foldable result: ", to_rev[0])
print("on result_logits: ", result_logits[0])
print("on mask: ", result_mask[0])
a = decode_one_ems_token_rec_for_folding_with_mask(
to_rev[0],
result_logits[0],
pLM_alphabet,
pLM_Model,
result_mask[0],
)
print('One resu: ', a)
y_data_reversed=decode_many_ems_token_rec_for_folding_with_mask(
to_rev,
result_logits,
pLM_alphabet,
pLM_Model,
result_mask,
)
if CKeys['Debug_TrainerPack']==3:
print("on y_data_reversed[0]: ", y_data_reversed[0])
# # ++++++++++++++++++++++++++++++++++++++++++++
# # reverse the PREDICTED y into a foldable sequence
# # save this block for Model A
# to_rev=result[:,0,:] # token (batch,seq_len)
# y_data_reversed=decode_many_ems_token_rec_for_folding(
# to_rev,
# result_logits,
# pLM_alphabet,
# pLM_Model,
# )
# if CKeys['Debug_TrainerPack']==3:
# print("on foldable result: ", to_rev[0])
# print("on result_logits: ", result_logits[0])
# a = decode_one_ems_token_rec_for_folding(
# to_rev[0],
# result_logits[0],
# pLM_alphabet,
# pLM_Model,
# )
# print('One resu: ', a)
# print("on y_data_reversed: ", y_data_reversed[0])
# # -----------------------------------------------------
# #reverse GT_y sequence
# to_rev=GT[:,0,:]
# to_rev=to_rev.long().cpu().detach().numpy()
#
# GT_y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
#
# for iii in range (len(y_data_reversed)):
# GT_y_data_reversed[iii]=GT_y_data_reversed[iii].upper().strip().replace(" ", "")
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
#reverse GT_y sequence
# GT should be SAFE to reverse
to_rev=GT[:,0,:] # (batch,1,seq_len)->(batch, seq_len)
GT_y_data_reversed=decode_many_ems_token_rec(
to_rev,
pLM_alphabet,
)
# -- not for SecStr anymore
# ### reverse second structure input....
# to_rev=torch.round (X_train_batch[:,:]*Xnormfac)
# to_rev=to_rev.long().cpu().detach().numpy()
# ++
### reverse general float input...
to_rev=X_train_batch[:,:]*Xnormfac
to_rev=to_rev.cpu().detach().numpy()
# here, assume X_train_batch is for cond_img: there are padding at both beginning and ending part
# so, first move the 0th padding to the end:
# Note:
# 1. this is good for SecStr case: (not maintained here)
# 2. this is not good for ForcPath, but can be cued in MD postprocess since the first component will always be 0
n_batch=to_rev.shape[0]
n_embed=to_rev.shape[1]
to_rev_1 = np.zeros(to_rev.shape)
to_rev_1[:,0:n_embed-1]=to_rev[:,1:n_embed]
# ++ different input
if tokenizer_X!=None:
# change into int
to_rev_1 = np.round(to_rev_1)
X_data_reversed=tokenizer_X.sequences_to_texts (to_rev_1)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
else:
X_data_reversed=to_rev_1.copy()
# + for debug
if CKeys['Debug_TrainerPack']==1:
print("X_data_reversed: ", X_data_reversed)
for samples in range (num_samples):
print ("sample ", samples+1, "out of ", num_samples)
fig=plt.figure()
plt.plot (
result[samples,0,:].cpu().detach().numpy(),
label= f'Predicted'
)
plt.plot (
GT[samples,0,:],
label= f'GT {0}'
)
plt.legend()
outname = sample_dir+ f"Batch_{idx}_sample_{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
if IF_showfig==1:
plt.show()
else:
plt.savefig(outname, dpi=200)
plt.close ()
# # # -------------------------------------------
# # #reverse y sequence
# # to_rev=result[:,0,:]
# # to_rev=to_rev.long().cpu().detach().numpy()
# #
# # y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
# #
# # for iii in range (len(y_data_reversed)):
# # y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
# # ++++++++++++++++++++++++++++++++++++++++++++
# # extract the mask/seq_len from input if possible
# # here, from dataloader, we only use tokenized_data for mask generation
# result_mask = read_mask_from_input(
# tokenized_data=X_train_batch[:num_samples],
# mask_value=0,
# seq_data=None,
# max_seq_length=None,
# )
# to_rev=result[:,0,:] # token (batch,seq_len)
# if CKeys['Debug_TrainerPack']==3:
# print("on foldable result: ", to_rev[0])
# print("on result_logits: ", result_logits[0])
# print("on mask: ", result_mask[0])
# a = decode_one_ems_token_rec_for_folding_with_mask(
# to_rev[0],
# result_logits[0],
# pLM_alphabet,
# pLM_Model,
# result_mask[0],
# )
# print('One resu: ', a)
# y_data_reversed=decode_many_ems_token_rec_for_folding_with_mask(
# to_rev,
# result_logits,
# pLM_alphabet,
# pLM_Model,
# result_mask,
# )
# if CKeys['Debug_TrainerPack']==3:
# print("on y_data_reversed[0]: ", y_data_reversed[0])
# # # ++++++++++++++++++++++++++++++++++++++++++++
# # # reverse the PREDICTED y into a foldable sequence
# # # save this block for Model A
# # to_rev=result[:,0,:] # token (batch,seq_len)
# # y_data_reversed=decode_many_ems_token_rec_for_folding(
# # to_rev,
# # result_logits,
# # pLM_alphabet,
# # pLM_Model,
# # )
# # if CKeys['Debug_TrainerPack']==3:
# # print("on foldable result: ", to_rev[0])
# # print("on result_logits: ", result_logits[0])
# # a = decode_one_ems_token_rec_for_folding(
# # to_rev[0],
# # result_logits[0],
# # pLM_alphabet,
# # pLM_Model,
# # )
# # print('One resu: ', a)
# # print("on y_data_reversed: ", y_data_reversed[0])
# # # -----------------------------------------------------
# # #reverse GT_y sequence
# # to_rev=GT[:,0,:]
# # to_rev=to_rev.long().cpu().detach().numpy()
# #
# # GT_y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
# #
# # for iii in range (len(y_data_reversed)):
# # GT_y_data_reversed[iii]=GT_y_data_reversed[iii].upper().strip().replace(" ", "")
# # ++++++++++++++++++++++++++++++++++++++++++++++++++++++
# #reverse GT_y sequence
# # GT should be SAFE to reverse
# to_rev=GT[:,0,:]
# GT_y_data_reversed=decode_many_ems_token_rec(
# to_rev,
# pLM_alphabet,
# )
# ### reverse second structure input....
# to_rev=torch.round (X_train_batch[:,:]*Xnormfac)
# to_rev=to_rev.long().cpu().detach().numpy()
# # here, assume X_train_batch is for cond_img: there are padding at both beginning and ending part
# # so, first move the 0th padding to the end
# n_batch=to_rev.shape[0]
# n_embed=to_rev.shape[1]
# to_rev_1 = np.zeros(to_rev.shape)
# to_rev_1[:,0:n_embed-1]=to_rev[:,1:n_embed]
# # ++ different input
# if tokenizer_X!=None:
# X_data_reversed=tokenizer_X.sequences_to_texts (to_rev_1)
# for iii in range (len(y_data_reversed)):
# X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
# else:
# X_data_reversed=to_rev_1.copy()
# # + for debug
# if CKeys['Debug_TrainerPack']==1:
# print("X_data_reversed: ", X_data_reversed)
# print (f"For {X_train_batch[samples,:].cpu().detach().numpy()} or {X_data_reversed[samples]}, \npredicted sequence: ", y_data_reversed[samples])
print (f"For {X_train_batch[samples,:].cpu().detach().numpy()} \nor\n {X_data_reversed[samples]}, ")
print (f"predicted sequence: {y_data_reversed[samples]}")
print (f"Ground truth: {GT_y_data_reversed[samples]}")
error=string_diff (y_data_reversed[samples], GT_y_data_reversed[samples])/len (GT_y_data_reversed[samples])
print(f"Recovery ratio(Ref): {1.-error}")
# move some
# # -- X_train_batch is normalized
# xbc=X_train_batch[samples,:].cpu().detach().numpy()
# # out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})
# out_nam_content=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})
# ++
xbc = X_data_reversed[samples]
if type(xbc)==str:
out_nam_content=xbc
else:
out_nam_content=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.4f" % xbc})
# 1. write out the input X in the dataloder
out_nam_inX=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.txt'
# + write the condition clearly
# X_data_reversed: an array
with open(out_nam_inX, "w") as inX_file:
# inX_file.write(f'{X_data_reversed[samples]}\n')
inX_file.write(out_nam_content)
# 2. write out the predictions
out_nam_OuY_PR=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}_predict.fasta'
with open(out_nam_OuY_PR, "w") as ouY_fasta:
ouY_fasta.write(f">Predicted\n")
ouY_fasta.write(y_data_reversed[samples])
# 3. Only for dataloader: write out the recovered ground truth
out_nam_OuY_GT=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}_recGT.fasta'
with open(out_nam_OuY_GT, "w") as ouY_fasta:
ouY_fasta.write(f">reconstructed GT, recoverabliblity: {1.-error}\n")
ouY_fasta.write(GT_y_data_reversed[samples])
if foldproteins:
tempname='temp'
pdb_file,fasta_file=foldandsavePDB_pdb_fasta (
sequence=y_data_reversed[samples],
filename_out=tempname,
num_cycle=16, flag=flag,
# +++++++++++++++++++
prefix=prefix
)
# #out_nam=f'{prefix}{out_nam}.pdb'
# out_nam=f'{prefix}{X_data_reversed[samples]}.pdb'
# ------------------------------------------------------
# sometime, this name below can get too long to fit
# out_nam=f'{sample_dir}{X_data_reversed[samples]}.pdb'
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
# add a way to save the sampling name and results
# ref: outname = sample_dir+ f"sample-{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
out_nam=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.pdb'
out_nam_seq=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.fasta'
if CKeys['Debug_TrainerPack']==1:
print("pdb_file: ", pdb_file)
print("out_nam: ", out_nam)
print (f'Original PDB: {pdb_file} OUT: {out_nam}')
shutil.copy (pdb_file, out_nam) #source, dest
shutil.copy (fasta_file, out_nam_seq)
# clean the slade to avoid mistakenly using the previous fasta file
os.remove (pdb_file)
os.remove (fasta_file)
pdb_file=out_nam
print (f"Properly named PDB file produced: {pdb_file}")
print (f"input X for sampling stored: {pdb_file}")
if IF_showfig==1:
view=show_pdb(
pdb_file=pdb_file,
flag=flag,
show_sidechains=show_sidechains,
show_mainchains=show_mainchains,
color=color
)
view.show()
# ++
# For ProteinPredictor
def sample_loop_omegafold_pLM_ModelB_Predictor (
model,
train_loader,
cond_scales=[7.5], #list of cond scales - each sampled...
num_samples=2, #how many samples produced every time tested.....
timesteps=100, # not used
flag=0,
foldproteins=False,
use_text_embedd=True,
skip_steps=0,
# +++++++++++++++++++
train_unet_number=1,
ynormfac=1,
prefix=None,
tokenizer_y=None,
Xnormfac=1,
tokenizer_X=None,
# ++
CKeys=None,
sample_dir=None,
steps=None,
e=None,
IF_showfig=True, # effective only after foldproteins=True
# ++
pLM_Model=None,
pLM_Model_Name=None,
image_channels=None,
pLM_alphabet=None,
# ++
esm_layer=None,
):
# =====================================================
# sample # = num_samples*(# of mini-batches)
# =====================================================
# steps=0
# e=flag
# for item in train_loader:
val_epoch_MSE_list=[]
resu_pred = {}
resu_grou = {}
#
for iisample in range (len (cond_scales)):
# calculate loss for one selected cond_scales
#
val_epoch_MSE=0.
num_rec=0
this_prediction = []
this_groundtruth = []
#
for idx, item in enumerate(train_loader):
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# 1. adjust the number of sample to collect in each batch
if num_samples>y_train_batch.shape[0]:
print("Warning: sampling # > len(mini_batch)")
num_samples = min (num_samples,y_train_batch.shape[0])
print (f"Producing {num_samples} samples...")
X_train_batch_picked = X_train_batch[:num_samples,:] # X_train_batch_in[:num_samples ] #
GT=y_train_batch.cpu().detach()
GT_picked = GT[:num_samples,:]
# GT_picked = GT_picked.unsqueeze(1)
# 2. prepare if pLM is used at the input end:
# this is done inised model.sample fun via x_data_tokenized channel
#
# 3. sample inside the loop of cond_scales
if use_text_embedd:
result_embedding=model.sample (
# x= X_train_batch,
x= X_train_batch_picked,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample],
device=device,
skip_steps=skip_steps,
# ++
pLM_Model_Name=pLM_Model_Name,
pLM_Model=pLM_Model,
pLM_alphabet=pLM_alphabet,
esm_layer=esm_layer,
)
else:
result_embedding=model.sample (
x= None,
# x_data_tokenized= X_train_batch,
x_data_tokenized= X_train_batch_picked, # dim=(batch, seq_len), will extend channels inside .sample(),
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample],
device=device,
skip_steps=skip_steps,
# ++
pLM_Model_Name=pLM_Model_Name,
pLM_Model=pLM_Model,
pLM_alphabet=pLM_alphabet,
esm_layer=esm_layer,
)
# result_embedding as image.dim: [batch, channels, seq_len]
#
# 4. translate prediction into something meaningful
# consider channel average and masking
# average across channels
result_embedding=torch.mean(result_embedding, 1) # (batch, seq_len)
# read mask from input: X_train_batch_picked (batch, seq_len)
# result_mask looks like, 0,1,1,...,1,0,0
# will fill 0th component be zero
result_mask = read_mask_from_input(
tokenized_data=X_train_batch[:num_samples],
mask_value=0.0,
seq_data=None,
max_seq_length=None,
)
# apply mask to result: keep true and zero all false
result = result_embedding*result_mask # (batch, seq_len)
result = result.cpu()
# result = result.unsqueeze(1) # (batch, 1, seq_len)
# this is ONLY the result from the model, not predictio yet
#
# 5. calculate loss
with torch.no_grad():
val_loss_MSE = criterion_MSE_sum(
result,
GT_picked,
)
val_epoch_MSE += val_loss_MSE.item()/GT_picked.shape[1]
num_rec += len(GT_picked)
#
# 6. convert into prediction
y_data_reversed = result*ynormfac
# prepare GT
GT_y_data_reversed = GT_picked*ynormfac
# accumulate the results
for ibat in range (GT_picked.shape[0]):
this_prediction.append (np.array( y_data_reversed[ibat,:].cpu() ))
this_groundtruth.append (np.array( GT_y_data_reversed[ibat,:].cpu() ))
#
# 5. reverse input to AA sequence... if needed
# TBA
# for one scal_cond
# summarize the loss
TestSet_MSE = val_epoch_MSE/num_rec
resu_pred[str(cond_scales[iisample])] = this_prediction
resu_grou[str(cond_scales[iisample])] = this_groundtruth
# store the MSE along cond_scales
val_epoch_MSE_list.append(TestSet_MSE)
return val_epoch_MSE_list, resu_pred, resu_grou
# ++++++++++++++++++++++++++++++++++++++++++++++++
def sample_loop_omegafold_ModelB (
model,
train_loader,
cond_scales=[7.5], #list of cond scales - each sampled...
num_samples=2, #how many samples produced every time tested.....
timesteps=100, # not used
flag=0,
foldproteins=False,
use_text_embedd=True,
skip_steps=0,
# +++++++++++++++++++
train_unet_number=1,
ynormfac=1,
prefix=None,
tokenizer_y=None,
Xnormfac=1,
tokenizer_X=None,
# ++
CKeys=None,
sample_dir=None,
steps=None,
e=None,
IF_showfig=True, # effective only after foldproteins=True
):
# =====================================================
# sample # = num_samples*(# of mini-batches)
# =====================================================
# steps=0
# e=flag
# for item in train_loader:
for idx, item in enumerate(train_loader):
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
GT=y_train_batch.cpu().detach()
GT= GT.unsqueeze(1)
if num_samples>y_train_batch.shape[0]:
print("Warning: sampling # > len(mini_batch)")
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
X_train_batch_picked = X_train_batch[:num_samples,:]
print ('(TEST) X_batch shape: ', X_train_batch_picked.shape)
for iisample in range (len (cond_scales)):
if use_text_embedd:
result=model.sample (
# x= X_train_batch,
x= X_train_batch_picked,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample],
device=device,
skip_steps=skip_steps
)
else:
result=model.sample (
x= None,
# x_data_tokenized= X_train_batch,
x_data_tokenized= X_train_batch_picked,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample],
device=device,
skip_steps=skip_steps
)
result=torch.round(result*ynormfac)
GT=torch.round (GT*ynormfac)
for samples in range (num_samples):
print ("sample ", samples+1, "out of ", num_samples)
fig=plt.figure()
plt.plot (
result[samples,0,:].cpu().detach().numpy(),
label= f'Predicted'
)
plt.plot (
GT[samples,0,:],
label= f'GT {0}'
)
plt.legend()
outname = sample_dir+ f"Batch_{idx}_sample_{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
if IF_showfig==1:
plt.show()
else:
plt.savefig(outname, dpi=200)
plt.close ()
#reverse y sequence
to_rev=result[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
#reverse GT_y sequence
to_rev=GT[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
GT_y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
GT_y_data_reversed[iii]=GT_y_data_reversed[iii].upper().strip().replace(" ", "")
### reverse second structure input....
to_rev=torch.round (X_train_batch[:,:]*Xnormfac)
to_rev=to_rev.long().cpu().detach().numpy()
# ++ different input
if tokenizer_X!=None:
X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
else:
X_data_reversed=to_rev.copy()
# + for debug
if CKeys['Debug_TrainerPack']==1:
print("X_data_reversed: ", X_data_reversed)
print (f"For {X_train_batch[samples,:].cpu().detach().numpy()} or {X_data_reversed[samples]}, \npredicted sequence: ", y_data_reversed[samples])
print (f"Ground truth: {GT_y_data_reversed[samples]}")
if foldproteins:
xbc=X_train_batch[samples,:].cpu().detach().numpy()
out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})
tempname='temp'
pdb_file=foldandsavePDB (
sequence=y_data_reversed[samples],
filename_out=tempname,
num_cycle=16, flag=flag,
# +++++++++++++++++++
prefix=prefix
)
# #out_nam=f'{prefix}{out_nam}.pdb'
# out_nam=f'{prefix}{X_data_reversed[samples]}.pdb'
# ------------------------------------------------------
# sometime, this name below can get too long to fit
# out_nam=f'{sample_dir}{X_data_reversed[samples]}.pdb'
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
# add a way to save the sampling name and results
# ref: outname = sample_dir+ f"sample-{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
out_nam=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.pdb'
out_nam_inX=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.txt'
if CKeys['Debug_TrainerPack']==1:
print("pdb_file: ", pdb_file)
print("out_nam: ", out_nam)
print (f'Original PDB: {pdb_file} OUT: {out_nam}')
shutil.copy (pdb_file, out_nam) #source, dest
# +
with open(out_nam_inX, "w") as inX_file:
inX_file.write(f'{X_data_reversed[samples]}\n')
pdb_file=out_nam
print (f"Properly named PDB file produced: {pdb_file}")
print (f"input X for sampling stored: {pdb_file}")
if IF_showfig==1:
view=show_pdb(
pdb_file=pdb_file,
flag=flag,
show_sidechains=show_sidechains,
show_mainchains=show_mainchains,
color=color
)
view.show()
# steps=steps+1
# if steps>num_samples:
# break
#
#
def sample_loop_FromModelB (model,
train_loader,
cond_scales=[7.5], #list of cond scales - each sampled...
num_samples=2, #how many samples produced every time tested.....
timesteps=100,
flag=0,foldproteins=False,
use_text_embedd=True,skip_steps=0,
# +++++++++++++++++++
train_unet_number=1,
ynormfac=1,
prefix=None,
tokenizer_y=None,
Xnormfac=1,
tokenizer_X=None,
):
steps=0
e=flag
for item in train_loader:
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
GT=y_train_batch.cpu().detach()
GT= GT.unsqueeze(1)
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
print ('X_train_batch shape: ', X_train_batch.shape)
for iisample in range (len (cond_scales)):
if use_text_embedd:
result=model.sample (x= X_train_batch,stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample], device=device, skip_steps=skip_steps)
else:
result=model.sample (x= None, x_data_tokenized= X_train_batch,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales[iisample],device=device,skip_steps=skip_steps)
result=torch.round(result*ynormfac)
GT=torch.round (GT*ynormfac)
for samples in range (num_samples):
print ("sample ", samples, "out of ", num_samples)
plt.plot (result[samples,0,:].cpu().detach().numpy(),label= f'Predicted')
plt.plot (GT[samples,0,:],label= f'GT {0}')
plt.legend()
outname = prefix+ f"sample-{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
plt.savefig(outname, dpi=200)
plt.show ()
#reverse y sequence
to_rev=result[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
#reverse GT_y sequence
to_rev=GT[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
GT_y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
GT_y_data_reversed[iii]=GT_y_data_reversed[iii].upper().strip().replace(" ", "")
### reverse second structure input....
to_rev=torch.round (X_train_batch[:,:]*Xnormfac)
to_rev=to_rev.long().cpu().detach().numpy()
X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
print (f"For {X_train_batch[samples,:].cpu().detach().numpy()} or {X_data_reversed[samples]}, predicted sequence: ", y_data_reversed[samples])
print (f"Ground truth: {GT_y_data_reversed[samples]}")
if foldproteins:
xbc=X_train_batch[samples,:].cpu().detach().numpy()
out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})
tempname='temp'
pdb_file=foldandsavePDB (
sequence=y_data_reversed[samples],
filename_out=tempname,
num_cycle=16, flag=flag,
# +++++++++++++++++++
prefix=prefix
)
#out_nam=f'{prefix}{out_nam}.pdb'
out_nam=f'{prefix}{X_data_reversed[samples]}.pdb'
print (f'Original PDB: {pdb_file} OUT: {out_nam}')
shutil.copy (pdb_file, out_nam) #source, dest
pdb_file=out_nam
print (f"Properly named PDB file produced: {pdb_file}")
view=show_pdb(pdb_file=pdb_file, flag=flag, show_sidechains=show_sidechains, show_mainchains=show_mainchains, color=color)
view.show()
steps=steps+1
if steps>num_samples:
break
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# train_loop tasks:
# 1. calculate loss for one batch
# 2. call sample loop
# 3. call sample sequence
# 4. print records and save model
# ===============================================================
# for ProteinDesigner_B
# 1. expanded for Probability case
# ++
cal_norm_prob = nn.Softmax(dim=2)
def train_loop_Model_B (
model,
train_loader,
test_loader,
#
optimizer=None,
print_every=10,
epochs= 300,
start_ep=0,
start_step=0,
train_unet_number=1,
print_loss_every_steps=1000,
#
trainer=None,
plot_unscaled=False,
max_batch_size=4,
save_model=False,
cond_scales=[7.5], #list of cond scales - each sampled...
num_samples=2, #how many samples produced every time tested.....
foldproteins=False,
cond_image=False, #use cond_images...
# add some
# +++++++++++++++++++++++++++
device=None,
loss_list=[],
epoch_list=[],
train_hist_file=None,
train_hist_file_full=None,
prefix=None,
Xnormfac=1.,
ynormfac=1.,
tokenizer_X=None,
tokenizer_y=None,
test_condition_list=[],
max_length=1,
CKeys=None,
sample_steps=1,
sample_dir=None,
save_every_epoch=1,
save_point_info_file=None,
store_dir=None,
# ++
pLM_Model_Name=None,
image_channels=None,
print_error=False,
):
if not exists (trainer):
if not exists (optimizer):
print ("ERROR: If trainer not used, need to provide optimizer.")
if exists (trainer):
print ("Trainer provided... will be used")
# steps=start_step+1
# # +
# added_steps=0+1
#
steps=start_step
added_steps=0
loss_total=0
# ++ for pLM
if pLM_Model_Name=='None':
pLM_Model=None
elif pLM_Model_Name=='esm2_t33_650M_UR50D':
# dim: 1280
esm_layer=33
pLM_Model, esm_alphabet = esm.pretrained.esm2_t33_650M_UR50D()
len_toks=len(esm_alphabet.all_toks)
pLM_Model.eval()
pLM_Model. to(device)
elif pLM_Model_Name=='esm2_t36_3B_UR50D':
# dim: 2560
esm_layer=36
pLM_Model, esm_alphabet = esm.pretrained.esm2_t36_3B_UR50D()
len_toks=len(esm_alphabet.all_toks)
pLM_Model.eval()
pLM_Model. to(device)
elif pLM_Model_Name=='esm2_t30_150M_UR50D':
# dim: 640
esm_layer=30
pLM_Model, esm_alphabet = esm.pretrained.esm2_t30_150M_UR50D()
len_toks=len(esm_alphabet.all_toks)
pLM_Model.eval()
pLM_Model. to(device)
elif pLM_Model_Name=='esm2_t12_35M_UR50D':
# dim: 480
esm_layer=12
pLM_Model, esm_alphabet = esm.pretrained.esm2_t12_35M_UR50D()
len_toks=len(esm_alphabet.all_toks)
pLM_Model.eval()
pLM_Model. to(device)
else:
print("pLM model is missing...")
for e in range(1, epochs+1):
# start = time.time()
torch.cuda.empty_cache()
print ("######################################################################################")
start = time.time()
print ("NOW: Training epoch: ", e+start_ep)
# TRAINING
train_epoch_loss = 0
model.train()
print ("Loop over ", len(train_loader), " batches (print . every ", print_every, " steps)")
for item in train_loader:
steps += 1
added_steps += 1
X_train_batch= item[0].to(device)
y_train_batch= item[1].to(device)
# project y_ into embedding space
if CKeys["Debug_TrainerPack"]==1:
print("Initial unload the dataloader items: ...")
print("X_train_batch.dim: ", X_train_batch.shape)
print("y_train_batch.dim: ", y_train_batch.shape)
# ---------------------------------------------------------
# prepare for model.forward() to calculate loss
# ---------------------------------------------------------
# # --
# if pLM_Model_Name=='None':
# # just use the encoded sequence
# y_train_batch_in = y_train_batch.unsqueeze(1)
# X_train_batch_in = X_train_batch.unsqueeze(1)
# # pass
# elif pLM_Model_Name=='esm2_t33_650M_UR50D':
# with torch.no_grad():
# results = pLM_Model(
# y_train_batch,
# repr_layers=[33],
# return_contacts=False,
# )
# y_train_batch_in = results["representations"][33] # (batch, seq_len, channels)
# y_train_batch_in = rearrange(
# y_train_batch_in,
# 'b l c -> b c l'
# )
# X_train_batch_in = X_train_batch.unsqueeze(1).repeat(1,image_channels,1)
# else:
# print(f"Required pLM name is not defined!!")
# ++
if pLM_Model_Name=='None':
# just use the encoded sequence
y_train_batch_in = y_train_batch.unsqueeze(1)
X_train_batch_in = X_train_batch.unsqueeze(1)
# pass
else: # assume ESM models
with torch.no_grad():
results = pLM_Model(
y_train_batch,
repr_layers=[esm_layer],
return_contacts=False,
)
y_train_batch_in = results["representations"][esm_layer] # (batch, seq_len, channels)
# ++ for Probability case
if image_channels==33:
with torch.no_grad():
# calculate logits: (batch, seq_len, 33)
y_train_batch_in = pLM_Model.lm_head(
y_train_batch_in
)
# normalize to get (0,1) probability
y_train_batch_in = cal_norm_prob(y_train_batch_in)
# switch the dimension -> (batch, channel, seq_len)
y_train_batch_in = rearrange(
y_train_batch_in,
'b l c -> b c l'
)
X_train_batch_in = X_train_batch.unsqueeze(1).repeat(1,image_channels,1)
# + for debug
if CKeys["Debug_TrainerPack"]==1:
print("After pLM model, the shape of X and y for training:")
print("X_train_batch_in.dim: ", X_train_batch_in.shape)
print("y_train_batch_in.dim: ", y_train_batch_in.shape)
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
if exists (trainer):
if cond_image==False:
loss = trainer(
y_train_batch.unsqueeze(1) , # true image (batch, channels, seq_len)
x=X_train_batch, # tokenized text (batch, )
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
# # ----------------------------------------------------------
# if cond_image==True:
# loss = trainer(
# y_train_batch.unsqueeze(1) , # true image
# x=None, # tokenized text
# cond_images=X_train_batch.unsqueeze(1), # cond_image
# unet_number=train_unet_number,
# max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
# )
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
if cond_image==True:
loss = trainer(
y_train_batch_in, # true image
x=None, # tokenized text
cond_images=X_train_batch_in, # cond_image
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
trainer.update(unet_number = train_unet_number)
else:
optimizer.zero_grad()
if cond_image==False:
loss=model (
y_train_batch.unsqueeze(1) ,
x=X_train_batch,
unet_number=train_unet_number
)
# # ------------------------------------------------------
# if cond_image==True:
# loss=model (
# y_train_batch.unsqueeze(1) ,
# x=None,
# cond_images=X_train_batch.unsqueeze(1),
# unet_number=train_unet_number
# )
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
if cond_image==True:
loss=model (
y_train_batch_in ,
x=None,
cond_images=X_train_batch_in,
unet_number=train_unet_number
)
#
loss.backward( )
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
loss_total=loss_total+loss.item()
# +
train_epoch_loss=train_epoch_loss+loss.item()
if steps % print_every == 0:
# for progress bar
print(".", end="")
# if steps>0:
if added_steps>0:
if steps % print_loss_every_steps == 0:
# + for debug
if CKeys['Debug_TrainerPack']==2:
print('I am here')
print("Here is steps: ", steps)
norm_loss=loss_total/print_loss_every_steps
print (f"\nTOTAL LOSS at epoch={e+start_ep}, step={steps}: {norm_loss}")
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# add a line to the hist file
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss)+'\n'
with open(train_hist_file,'a') as f:
f.write(add_line)
loss_list.append (norm_loss)
loss_total=0
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
epoch_list.append(e+start_ep)
# fig = plt.figure(figsize=(12,8),dpi=200)
fig = plt.figure()
plt.plot (epoch_list, loss_list, label='Loss')
plt.legend()
# outname = prefix+ f"loss_{e+start_ep}_{steps}.jpg"
outname = sample_dir+ f"loss_{e+start_ep}_{steps}.jpg"
#
# the order, save then show, matters
if CKeys['SlientRun']==1:
plt.savefig(outname, dpi=200)
else:
plt.show()
plt.close(fig)
# plt.close()
if added_steps>0:
# if steps>0:
if steps % sample_steps == 0:
# + for debug
if CKeys['Debug_TrainerPack']==2:
print('I am here')
print("Here is steps: ", steps)
if plot_unscaled:
#test before scaling...
plt.plot (
y_train_batch.unsqueeze(1)[0,0,:].cpu().detach().numpy(),
label= 'Unscaled GT'
)
plt.legend()
plt.show()
# # --------------------------------------------------
# GT=y_train_batch.cpu().detach()
# GT=resize_image_to(
# GT.unsqueeze(1),
# model.imagen.image_sizes[train_unet_number-1],
# )
####
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
print ("I. SAMPLING IN TEST SET: ")
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
####
# num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
if cond_image == True:
use_text_embedd=False
# -
# cond_scales_extended=[1. for i in range(num_samples)]
# +
cond_scales_extended=cond_scales
else:
use_text_embedd=True
cond_scales_extended=cond_scales
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
sample_loop_omegafold_pLM_ModelB (
model,
test_loader,
cond_scales=cond_scales_extended, # cond_scales,# #list of cond scales - each sampled...
num_samples=num_samples, #how many samples produced every time tested.....
timesteps=64,
flag=steps,
#reverse=False,
foldproteins=foldproteins,
use_text_embedd= use_text_embedd,
# ++++++++++++++++++++
train_unet_number=train_unet_number,
ynormfac=ynormfac,
prefix=prefix,
tokenizer_y=tokenizer_y,
Xnormfac=Xnormfac,
tokenizer_X=tokenizer_X,
# ++
# ++
CKeys=CKeys,
sample_dir=sample_dir,
steps=steps,
e=e+start_ep,
IF_showfig= CKeys['SlientRun']!=1,
# ++
pLM_Model=pLM_Model,
pLM_Model_Name=pLM_Model_Name,
image_channels=image_channels,
pLM_alphabet=esm_alphabet,
)
#---------------------------------------------------------
# sample_loop (
# model,
# test_loader,
# cond_scales=cond_scales,# #list of cond scales - each sampled...
# num_samples=num_samples, #how many samples produced every time tested.....
# timesteps=64,
# flag=steps,
# #reverse=False,
# foldproteins=foldproteins,
# use_text_embedd= use_text_embedd,
# # ++++++++++++++++++++
# train_unet_number=train_unet_number,
# ynormfac=ynormfac,
# prefix=prefix,
# tokenizer_y=tokenizer_y,
# Xnormfac=Xnormfac,
# tokenizer_X=tokenizer_X,
# )
#index_word': '{"1": "~", "2": "h", "3": "e", "4": "s", "5": "t", "6": "g", "7": "b", "8": "i"}',
#'word_index': '{"~": 1, "h": 2, "e": 3, "s": 4, "t": 5, "g": 6, "b": 7, "i": 8}'}
AH_code=2/Xnormfac
BS_code=3/Xnormfac
unstr_code= 1/Xnormfac
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
print ("II. SAMPLING FOR DE NOVO:")
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
# +++++++++++++++++++++++++++++++++++++++++
DeNovoSam_pdbs, fasta_file_list=\
sample_sequence_omegafold_pLM_ModelB (
model,
x_data=test_condition_list,
flag=steps, # flag="DeNovo", # ,
cond_scales=1.,
foldproteins=foldproteins,
# ++++++++++
ynormfac=ynormfac,
train_unet_number=train_unet_number,
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
prefix=prefix,
tokenizer_y=tokenizer_y,
# ++
CKeys=CKeys,
sample_dir=sample_dir,
steps=steps,
e=e+start_ep,
IF_showfig= CKeys['SlientRun']!=1,
# ++
pLM_Model=pLM_Model,
pLM_Model_Name=pLM_Model_Name,
image_channels=image_channels,
pLM_alphabet=esm_alphabet,
)
if print_error and len(DeNovoSam_pdbs)>0:
print("Calculate SecStr and design error:")
#
for ii in range(len(test_condition_list)):
seq=test_condition_list[ii][0]
DSSPresult,_,sequence_res=get_DSSP_result(DeNovoSam_pdbs[ii])
print (f"INPUT: {seq}\nRESULT: {DSSPresult}\nAA sequence: {sequence_res}")
error=string_diff (DSSPresult, seq)/len (seq)
print ("Error: ", error)
# # +++++++++++++++++++++++++++++++++++++++++
# sample_sequence_omegafold_ModelB (
# model,
# x_data=test_condition_list,
# flag=steps, # flag="DeNovo", # ,
# cond_scales=1.,
# foldproteins=foldproteins,
# # ++++++++++
# ynormfac=ynormfac,
# train_unet_number=train_unet_number,
# tokenizer_X=tokenizer_X,
# Xnormfac=Xnormfac,
# max_length=max_length,
# prefix=prefix,
# tokenizer_y=tokenizer_y,
# # ++
# CKeys=CKeys,
# sample_dir=sample_dir,
# steps=steps,
# e=e+start_ep,
# IF_showfig= CKeys['SlientRun']!=1,
# )
# for this_x_data in test_condition_list:
# sample_sequence_omegafold (
# model,
# x_data=this_x_data,
# flag=steps, # flag="DeNovo", # ,
# cond_scales=1.,
# foldproteins=True,
# # ++++++++++
# ynormfac=ynormfac,
# train_unet_number=train_unet_number,
# tokenizer_X=tokenizer_X,
# Xnormfac=Xnormfac,
# max_length=max_length,
# prefix=prefix,
# tokenizer_y=tokenizer_y,
# # ++
# CKeys=CKeys,
# sample_dir=sample_dir,
# steps=steps,
# e=e+start_ep,
# IF_showfig= CKeys['SlientRun']!=1,
# )
#
# model,
# X=None, #this is the target conventionally when using text embd
# flag=0,
# cond_scales=1.,
# foldproteins=False,
# X_string=None,
# x_data=None,
# skip_steps=0,
# inpaint_images = None,
# inpaint_masks = None,
# inpaint_resample_times = None,
# init_images = None,
# num_cycle=16,
# # ++++++++++++++++++++++++
# ynormfac=1,
# train_unet_number=1,
# tokenizer_X=None,
# Xnormfac=1.,
# max_length=1.,
# prefix=None,
# tokenizer_y=None,
# # ++
# CKeys=None,
# sample_dir=None,
# -----------------------------------------
# sample_sequence (
# model,
# x_data=['~~~HHHHHHHHHHHHHHH~~'],
# flag=steps,cond_scales=1.,
# foldproteins=True,
# # ++++++++++
# ynormfac=ynormfac,
# )
# sample_sequence (
# model,
# x_data=['~~~HHHHHHHHHHHHHHH~~~~HHHHHHHHHHHHHH~~~'],
# flag=steps,cond_scales=1.,
# foldproteins=True,
# # ++++++++++
# ynormfac=ynormfac,
# )
# sample_sequence (
# model,
# x_data=['~~EEESSTTS~SEEEEEEEEE~SBS~EEEEEE~~'],
# flag=steps,cond_scales=1.,
# foldproteins=True,
# # ++++++++++++
# ynormfac=ynormfac,
# )
# if steps>0:
# # --------------------------------------------------------------------
# if added_steps>0:
# if save_model and steps % print_loss_every_steps==0:
# fname=f"{prefix}trainer_save-model-epoch_{e+start_ep}.pt"
# trainer.save(fname)
# print (f"Model saved: ", fname)
# fname=f"{prefix}statedict_save-model-epoch_{e+start_ep}.pt"
# torch.save(model.state_dict(), fname)
# print (f"Statedict model saved: ", fname)
# steps=steps+1
# added_steps += 1
# every epoch:
norm_loss_over_e = train_epoch_loss/len(train_loader)
print("\nnorm_loss over 1 epoch: ", norm_loss_over_e)
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
# write this into "train_hist_file_full"
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss_over_e)+'\n'
with open(train_hist_file_full,'a') as f:
f.write(add_line)
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# save model every this epoches
if save_model and (e+start_ep) % save_every_epoch==0 and e>1:
# fname=f"{prefix}trainer_save-model-epoch_{e+start_ep}.pt"
fname=f"{store_dir}trainer_save-model-epoch_{e+start_ep}.pt"
trainer.save(fname)
print (f"Model saved: ", fname)
# fname=f"{prefix}statedict_save-model-epoch_{e+start_ep}.pt"
fname=f"{store_dir}statedict_save-model-epoch_{e+start_ep}.pt"
torch.save(model.state_dict(), fname)
print (f"Statedict model saved: ", fname)
# add a saving point file
top_line='epoch,steps,norm_loss'+'\n'
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss)+'\n'
with open(save_point_info_file, "w") as f:
f.write(top_line)
f.write(add_line)
print (f"\n\n-------------------\nTime for epoch {e+start_ep}={(time.time()-start)/60}\n-------------------")
# ===============================================================
# for ProteinPredictor_B
#
def train_loop_Model_B_Predictor (
model,
train_loader,
test_loader,
#
optimizer=None,
print_every=10,
epochs= 300,
start_ep=0,
start_step=0,
train_unet_number=1,
print_loss_every_steps=1000,
#
trainer=None,
plot_unscaled=False,
max_batch_size=4,
save_model=False,
cond_scales=[1.], #list of cond scales - each sampled...
num_samples=2, #how many samples produced every time tested.....
foldproteins=False,
cond_image=False, #use cond_images...
# add some
# +++++++++++++++++++++++++++
device=None,
loss_list=[],
epoch_list=[],
train_hist_file=None,
train_hist_file_full=None,
prefix=None,
Xnormfac=1.,
ynormfac=1.,
tokenizer_X=None,
tokenizer_y=None,
test_condition_list=[],
max_length=1,
CKeys=None,
sample_steps=1,
sample_dir=None,
save_every_epoch=1,
save_point_info_file=None,
store_dir=None,
# ++
pLM_Model_Name=None,
image_channels=None,
print_error=False,
# ++
train_hist_file_on_testset=None,
):
if not exists (trainer):
if not exists (optimizer):
print ("ERROR: If trainer not used, need to provide optimizer.")
if exists (trainer):
print ("Trainer provided... will be used")
# steps=start_step+1
# # +
# added_steps=0+1
#
steps=start_step
added_steps=0
loss_total=0
# ++ for pLM
# # --
# if pLM_Model_Name=='trivial':
# pLM_Model=None
# elif pLM_Model_Name=='esm2_t33_650M_UR50D':
# # dim: 1280
# esm_layer=33
# pLM_Model, esm_alphabet = esm.pretrained.esm2_t33_650M_UR50D()
# len_toks=len(esm_alphabet.all_toks)
# pLM_Model.eval()
# pLM_Model. to(device)
# elif pLM_Model_Name=='esm2_t36_3B_UR50D':
# # dim: 2560
# esm_layer=36
# pLM_Model, esm_alphabet = esm.pretrained.esm2_t36_3B_UR50D()
# len_toks=len(esm_alphabet.all_toks)
# pLM_Model.eval()
# pLM_Model. to(device)
# elif pLM_Model_Name=='esm2_t30_150M_UR50D':
# # dim: 640
# esm_layer=30
# pLM_Model, esm_alphabet = esm.pretrained.esm2_t30_150M_UR50D()
# len_toks=len(esm_alphabet.all_toks)
# pLM_Model.eval()
# pLM_Model. to(device)
# elif pLM_Model_Name=='esm2_t12_35M_UR50D':
# # dim: 480
# esm_layer=12
# pLM_Model, esm_alphabet = esm.pretrained.esm2_t12_35M_UR50D()
# len_toks=len(esm_alphabet.all_toks)
# pLM_Model.eval()
# pLM_Model. to(device)
# else:
# print("pLM model is missing...")
# ++
pLM_Model, esm_alphabet, \
esm_layer, len_toks = load_in_pLM(
pLM_Model_Name,
device,
)
for e in range(1, epochs+1):
# start = time.time()
torch.cuda.empty_cache()
print ("######################################################################################")
start = time.time()
print ("NOW: Training epoch: ", e+start_ep)
# TRAINING
train_epoch_loss = 0
model.train()
print ("Loop over ", len(train_loader), " batches (print . every ", print_every, " steps)")
for item in train_loader:
steps += 1
added_steps += 1
X_train_batch= item[0].to(device)
y_train_batch= item[1].to(device)
# project y_ into embedding space
if CKeys["Debug_TrainerPack"]==1:
print("Initial unload the dataloader items: ...")
print("X_train_batch.dim: ", X_train_batch.shape)
print("y_train_batch.dim: ", y_train_batch.shape)
# ---------------------------------------------------------
# prepare for model.forward() to calculate loss
# ---------------------------------------------------------
# # --
# if pLM_Model_Name=='None':
# # just use the encoded sequence
# y_train_batch_in = y_train_batch.unsqueeze(1)
# X_train_batch_in = X_train_batch.unsqueeze(1)
# # pass
# elif pLM_Model_Name=='esm2_t33_650M_UR50D':
# with torch.no_grad():
# results = pLM_Model(
# y_train_batch,
# repr_layers=[33],
# return_contacts=False,
# )
# y_train_batch_in = results["representations"][33] # (batch, seq_len, channels)
# y_train_batch_in = rearrange(
# y_train_batch_in,
# 'b l c -> b c l'
# )
# X_train_batch_in = X_train_batch.unsqueeze(1).repeat(1,image_channels,1)
# else:
# print(f"Required pLM name is not defined!!")
# ++
if pLM_Model_Name=='trivial':
# just use the encoded sequence
y_train_batch_in = y_train_batch.unsqueeze(1)
X_train_batch_in = X_train_batch.unsqueeze(1)
# pass
else:
# assume ESM models
# --
# # for ProteinDesigner
# with torch.no_grad():
# results = pLM_Model(
# y_train_batch,
# repr_layers=[esm_layer],
# return_contacts=False,
# )
# y_train_batch_in = results["representations"][esm_layer] # (batch, seq_len, channels)
# y_train_batch_in = rearrange(
# y_train_batch_in,
# 'b l c -> b c l'
# )
# X_train_batch_in = X_train_batch.unsqueeze(1).repeat(1,image_channels,1)
#
# ++
# for ProteinPredictor
with torch.no_grad():
results = pLM_Model(
X_train_batch,
repr_layers=[esm_layer],
return_contacts=False,
)
X_train_batch_in = results["representations"][esm_layer] # (batch, seq_len, channels)
X_train_batch_in = rearrange(
X_train_batch_in,
'b l c -> b c l'
)
y_train_batch_in = y_train_batch.unsqueeze(1).repeat(1,image_channels,1)
# + for debug
if CKeys["Debug_TrainerPack"]==1:
print("After pLM model, the shape of X and y for training:")
print("X_train_batch_in.dim: ", X_train_batch_in.shape)
print("y_train_batch_in.dim: ", y_train_batch_in.shape)
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
if exists (trainer):
if cond_image==False:
loss = trainer(
y_train_batch.unsqueeze(1) , # true image (batch, channels, seq_len)
x=X_train_batch, # tokenized text (batch, )
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
# # ----------------------------------------------------------
# if cond_image==True:
# loss = trainer(
# y_train_batch.unsqueeze(1) , # true image
# x=None, # tokenized text
# cond_images=X_train_batch.unsqueeze(1), # cond_image
# unet_number=train_unet_number,
# max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
# )
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
if cond_image==True:
loss = trainer(
y_train_batch_in, # true image
x=None, # tokenized text
cond_images=X_train_batch_in, # cond_image
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
trainer.update(unet_number = train_unet_number)
else:
optimizer.zero_grad()
if cond_image==False:
loss=model (
y_train_batch.unsqueeze(1) ,
x=X_train_batch,
unet_number=train_unet_number
)
# # ------------------------------------------------------
# if cond_image==True:
# loss=model (
# y_train_batch.unsqueeze(1) ,
# x=None,
# cond_images=X_train_batch.unsqueeze(1),
# unet_number=train_unet_number
# )
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
if cond_image==True:
loss=model (
y_train_batch_in ,
x=None,
cond_images=X_train_batch_in,
unet_number=train_unet_number
)
#
loss.backward( )
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
loss_total=loss_total+loss.item()
# +
train_epoch_loss=train_epoch_loss+loss.item()
if steps % print_every == 0:
# for progress bar
print(".", end="")
# if steps>0:
if added_steps>0:
if steps % print_loss_every_steps == 0:
# + for debug
if CKeys['Debug_TrainerPack']==2:
print('I am here')
print("Here is steps: ", steps)
norm_loss=loss_total/print_loss_every_steps
print (f"\nTOTAL LOSS at epoch={e+start_ep}, step={steps}: {norm_loss}")
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# add a line to the hist file
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss)+'\n'
with open(train_hist_file,'a') as f:
f.write(add_line)
loss_list.append (norm_loss)
loss_total=0
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
epoch_list.append(e+start_ep)
# fig = plt.figure(figsize=(12,8),dpi=200)
fig = plt.figure()
plt.plot (epoch_list, loss_list, label='Loss')
plt.legend()
# outname = prefix+ f"loss_{e+start_ep}_{steps}.jpg"
outname = sample_dir+ f"loss_{e+start_ep}_{steps}.jpg"
#
# the order, save then show, matters
if CKeys['SlientRun']==1:
plt.savefig(outname, dpi=200)
else:
plt.show()
plt.close(fig)
# plt.close()
if added_steps>0:
# if steps>0:
if steps % sample_steps == 0:
# + for debug
if CKeys['Debug_TrainerPack']==2:
print('I am here')
print("Here is steps: ", steps)
if plot_unscaled:
#test before scaling...
plt.plot (
y_train_batch.unsqueeze(1)[0,0,:].cpu().detach().numpy(),
label= 'Unscaled GT'
)
plt.legend()
plt.show()
# # --------------------------------------------------
# GT=y_train_batch.cpu().detach()
# GT=resize_image_to(
# GT.unsqueeze(1),
# model.imagen.image_sizes[train_unet_number-1],
# )
####
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
print ("I. SAMPLING IN TEST SET: ")
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
####
# num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
if cond_image == True:
use_text_embedd=False
# -
# cond_scales_extended=[1. for i in range(num_samples)]
# +
cond_scales_extended=cond_scales
else:
use_text_embedd=True
cond_scales_extended=cond_scales
# # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# # For ProteinDesigner
# sample_loop_omegafold_pLM_ModelB (
# model,
# test_loader,
# cond_scales=cond_scales_extended, # cond_scales,# #list of cond scales - each sampled...
# num_samples=num_samples, #how many samples produced every time tested.....
# timesteps=64,
# flag=steps,
# #reverse=False,
# foldproteins=foldproteins,
# use_text_embedd= use_text_embedd,
# # ++++++++++++++++++++
# train_unet_number=train_unet_number,
# ynormfac=ynormfac,
# prefix=prefix,
# tokenizer_y=tokenizer_y,
# Xnormfac=Xnormfac,
# tokenizer_X=tokenizer_X,
# # ++
# # ++
# CKeys=CKeys,
# sample_dir=sample_dir,
# steps=steps,
# e=e+start_ep,
# IF_showfig= CKeys['SlientRun']!=1,
# # ++
# pLM_Model=pLM_Model,
# pLM_Model_Name=pLM_Model_Name,
# image_channels=image_channels,
# pLM_alphabet=esm_alphabet,
# )
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# For ProteinPredictor:
val_epoch_MSE_list, \
resu_pred, resu_grou = \
sample_loop_omegafold_pLM_ModelB_Predictor (
model,
test_loader,
cond_scales=[1.], # cond_scales_extended, # #list of cond scales - each sampled...
num_samples=num_samples, #how many samples produced every time tested.....
timesteps=64,
flag=steps,
#reverse=False,
foldproteins=foldproteins,
use_text_embedd= use_text_embedd,
# ++++++++++++++++++++
train_unet_number=train_unet_number,
ynormfac=ynormfac,
prefix=prefix,
tokenizer_y=tokenizer_y,
Xnormfac=Xnormfac,
tokenizer_X=tokenizer_X,
# ++
# ++
CKeys=CKeys,
sample_dir=sample_dir,
steps=steps,
e=e+start_ep,
IF_showfig= CKeys['SlientRun']!=1,
# ++
pLM_Model=pLM_Model,
pLM_Model_Name=pLM_Model_Name,
image_channels=image_channels,
pLM_alphabet=esm_alphabet,
# ++
esm_layer=esm_layer,
)
# record the ERROR on the test set
print(f"Epo {str(e+start_ep)}, on TestSet, MSE: {val_epoch_MSE_list[0]}")
# only write the 0th case of MSE
add_line = str(e+start_ep)+','+str(steps)+','+str(val_epoch_MSE_list[0])+'\n'
with open(train_hist_file_on_testset,'a') as f:
f.write(add_line)
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
print ("II. SAMPLING FOR DE NOVO: NOT USED in predictor mode")
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
# # +++++++++++++++++++++++++++++++++++++++++
# DeNovoSam_pdbs, fasta_file_list=\
# sample_sequence_omegafold_pLM_ModelB (
# model,
# x_data=test_condition_list,
# flag=steps, # flag="DeNovo", # ,
# cond_scales=1.,
# foldproteins=foldproteins,
# # ++++++++++
# ynormfac=ynormfac,
# train_unet_number=train_unet_number,
# tokenizer_X=tokenizer_X,
# Xnormfac=Xnormfac,
# max_length=max_length,
# prefix=prefix,
# tokenizer_y=tokenizer_y,
# # ++
# CKeys=CKeys,
# sample_dir=sample_dir,
# steps=steps,
# e=e+start_ep,
# IF_showfig= CKeys['SlientRun']!=1,
# # ++
# pLM_Model=pLM_Model,
# pLM_Model_Name=pLM_Model_Name,
# image_channels=image_channels,
# pLM_alphabet=esm_alphabet,
# )
# if print_error and len(DeNovoSam_pdbs)>0:
# print("Calculate SecStr and design error:")
# #
# for ii in range(len(test_condition_list)):
# seq=test_condition_list[ii][0]
# DSSPresult,_,sequence_res=get_DSSP_result(DeNovoSam_pdbs[ii])
# print (f"INPUT: {seq}\nRESULT: {DSSPresult}\nAA sequence: {sequence_res}")
# error=string_diff (DSSPresult, seq)/len (seq)
# print ("Error: ", error)
# every epoch:
norm_loss_over_e = train_epoch_loss/len(train_loader)
print("\nnorm_loss over 1 epoch: ", norm_loss_over_e)
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
# write this into "train_hist_file_full"
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss_over_e)+'\n'
with open(train_hist_file_full,'a') as f:
f.write(add_line)
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# save model every this epoches
if save_model and (e+start_ep) % save_every_epoch==0 and e>1:
# fname=f"{prefix}trainer_save-model-epoch_{e+start_ep}.pt"
fname=f"{store_dir}trainer_save-model-epoch_{e+start_ep}.pt"
trainer.save(fname)
print (f"Model saved: ", fname)
# fname=f"{prefix}statedict_save-model-epoch_{e+start_ep}.pt"
fname=f"{store_dir}statedict_save-model-epoch_{e+start_ep}.pt"
torch.save(model.state_dict(), fname)
print (f"Statedict model saved: ", fname)
# add a saving point file
top_line='epoch,steps,norm_loss'+'\n'
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss)+'\n'
with open(save_point_info_file, "w") as f:
f.write(top_line)
f.write(add_line)
print (f"\n\n-------------------\nTime for epoch {e+start_ep}={(time.time()-start)/60}\n-------------------")
def train_loop_Old_FromModelB (model,
train_loader,
test_loader,
optimizer=None,
print_every=10,
epochs= 300,
start_ep=0,
start_step=0,
train_unet_number=1,
print_loss=1000,
trainer=None,
plot_unscaled=False,
max_batch_size=4,
save_model=False,
cond_scales=[7.5], #list of cond scales - each sampled...
num_samples=2, #how many samples produced every time tested.....
foldproteins=False,
cond_image=False, #use cond_images...
# add some
# +++++++++++++++++++++++++++
device=None,
loss_list=[],
prefix=None,
ynormfac=1,
test_condition_list=[],
tokenizer_y=None,
Xnormfac=1,
tokenizer_X=None,
max_length=1,
):
if not exists (trainer):
if not exists (optimizer):
print ("ERROR: If trainer not used, need to provide optimizer.")
if exists (trainer):
print ("Trainer provided... will be used")
steps=start_step
loss_total=0
for e in range(1, epochs+1):
start = time.time()
torch.cuda.empty_cache()
print ("######################################################################################")
start = time.time()
print ("NOW: Training epoch: ", e+start_ep)
train_epoch_loss = 0
model.train()
print ("Loop over ", len(train_loader), " batches (print . every ", print_every, " steps)")
for item in train_loader:
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
if exists (trainer):
if cond_image==False:
loss = trainer(
y_train_batch.unsqueeze(1) ,
x=X_train_batch,
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
if cond_image==True:
loss = trainer(
y_train_batch.unsqueeze(1) ,x=None,
cond_images=X_train_batch.unsqueeze(1),
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
trainer.update(unet_number = train_unet_number)
else:
optimizer.zero_grad()
if cond_image==False:
loss=model (y_train_batch.unsqueeze(1) , x=X_train_batch, unet_number=train_unet_number)
if cond_image==True:
loss=model (y_train_batch.unsqueeze(1) ,x=None, cond_images=X_train_batch.unsqueeze(1), unet_number=train_unet_number)
loss.backward( )
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
loss_total=loss_total+loss.item()
if steps % print_every == 0:
print(".", end="")
if steps>0:
if steps % print_loss == 0:
if plot_unscaled:
#test before scaling...
plt.plot (y_train_batch.unsqueeze(1)[0,0,:].cpu().detach().numpy(),label= 'Unscaled GT')
plt.legend()
plt.show()
GT=y_train_batch.cpu().detach()
GT=resize_image_to(
GT.unsqueeze(1),
model.imagen.image_sizes[train_unet_number-1],
)
norm_loss=loss_total/print_loss
print (f"\nTOTAL LOSS at epoch={e}, step={steps}: {norm_loss}")
loss_list.append (norm_loss)
loss_total=0
plt.plot (loss_list, label='Loss')
plt.legend()
outname = prefix+ f"loss_{e}_{steps}.jpg"
plt.savefig(outname, dpi=200)
plt.show()
####
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
if cond_image == True:
use_text_embedd=False
else:
use_text_embedd=True
sample_loop_FromModelB (
model,
test_loader,
cond_scales=cond_scales,# #list of cond scales - each sampled...
num_samples=num_samples, #how many samples produced every time tested.....
timesteps=64,
flag=steps,
#reverse=False,
foldproteins=foldproteins,
use_text_embedd= use_text_embedd,
# ++++++++++++++++++++
train_unet_number=train_unet_number,
ynormfac=ynormfac,
prefix=prefix,
tokenizer_y=tokenizer_y,
Xnormfac=Xnormfac,
tokenizer_X=tokenizer_X,
)
#index_word': '{"1": "~", "2": "h", "3": "e", "4": "s", "5": "t", "6": "g", "7": "b", "8": "i"}',
#'word_index': '{"~": 1, "h": 2, "e": 3, "s": 4, "t": 5, "g": 6, "b": 7, "i": 8}'}
AH_code=2/Xnormfac
BS_code=3/Xnormfac
unstr_code= 1/Xnormfac
print ("SAMPLING FOR DE NOVO:")
# +++++++++++++++++++++++++++++++++++++++++
for this_x_data in test_condition_list:
sample_sequence_FromModelB (
model,
x_data=this_x_data,
flag=steps,cond_scales=1.,
foldproteins=True,
# ++++++++++
ynormfac=ynormfac,
train_unet_number=train_unet_number,
tokenizer_X=tokenizer_X,
Xnormfac=Xnormfac,
max_length=max_length,
prefix=prefix,
tokenizer_y=tokenizer_y,
)
# -----------------------------------------
# sample_sequence (
# model,
# x_data=['~~~HHHHHHHHHHHHHHH~~'],
# flag=steps,cond_scales=1.,
# foldproteins=True,
# # ++++++++++
# ynormfac=ynormfac,
# )
# sample_sequence (
# model,
# x_data=['~~~HHHHHHHHHHHHHHH~~~~HHHHHHHHHHHHHH~~~'],
# flag=steps,cond_scales=1.,
# foldproteins=True,
# # ++++++++++
# ynormfac=ynormfac,
# )
# sample_sequence (
# model,
# x_data=['~~EEESSTTS~SEEEEEEEEE~SBS~EEEEEE~~'],
# flag=steps,cond_scales=1.,
# foldproteins=True,
# # ++++++++++++
# ynormfac=ynormfac,
# )
if steps>0:
if save_model and steps % print_loss==0:
fname=f"{prefix}trainer_save-model-epoch_{e}.pt"
trainer.save(fname)
print (f"Model saved: ", fname)
fname=f"{prefix}statedict_save-model-epoch_{e}.pt"
torch.save(model.state_dict(), fname)
print (f"Statedict model saved: ", fname)
steps=steps+1
print (f"\n\n-------------------\nTime for epoch {e}={(time.time()-start)/60}\n-------------------")
# ++++++++++++++++++++++++++++++++++++++++++++++
def foldandsavePDB_pdb_fasta (
sequence,
filename_out,
num_cycle=16,
flag=0,
# ++++++++++++
prefix=None,
):
filename=f"{prefix}fasta_in_{flag}.fasta"
print ("Writing FASTA file: ", filename)
OUTFILE=f"{filename_out}_{flag}"
with open (filename, mode ='w') as f:
f.write (f'>{OUTFILE}\n')
f.write (f'{sequence}')
print (f"Now run OmegaFold.... on device={device}")
# !omegafold $filename $prefix --num_cycle $num_cycle --device=$device
cmd_line=F"omegafold {filename} {prefix} --num_cycle {num_cycle} --device={device}"
print(os.popen(cmd_line).read())
print ("Done OmegaFold")
# PDB_result=f"{prefix}{OUTFILE}.PDB"
PDB_result=f"{prefix}{OUTFILE}.pdb"
print (f"Resulting PDB file...: {PDB_result}")
return PDB_result, filename
def foldandsavePDB (
sequence,
filename_out,
num_cycle=16,
flag=0,
# ++++++++++++
prefix=None,
):
filename=f"{prefix}fasta_in_{flag}.fasta"
print ("Writing FASTA file: ", filename)
OUTFILE=f"{filename_out}_{flag}"
with open (filename, mode ='w') as f:
f.write (f'>{OUTFILE}\n')
f.write (f'{sequence}')
print (f"Now run OmegaFold.... on device={device}")
# !omegafold $filename $prefix --num_cycle $num_cycle --device=$device
cmd_line=F"omegafold {filename} {prefix} --num_cycle {num_cycle} --device={device}"
print(os.popen(cmd_line).read())
print ("Done OmegaFold")
# PDB_result=f"{prefix}{OUTFILE}.PDB"
PDB_result=f"{prefix}{OUTFILE}.pdb"
print (f"Resulting PDB file...: {PDB_result}")
return PDB_result
import py3Dmol
def plot_plddt_legend(dpi=100):
thresh = ['plDDT:','Very low (<50)','Low (60)','OK (70)','Confident (80)','Very high (>90)']
plt.figure(figsize=(1,0.1),dpi=dpi)
########################################
for c in ["#FFFFFF","#FF0000","#FFFF00","#00FF00","#00FFFF","#0000FF"]:
plt.bar(0, 0, color=c)
plt.legend(thresh, frameon=False,
loc='center', ncol=6,
handletextpad=1,
columnspacing=1,
markerscale=0.5,)
plt.axis(False)
return plt
color = "lDDT" # choose from ["chain", "lDDT", "rainbow"]
show_sidechains = False #choose from {type:"boolean"}
show_mainchains = False #choose from {type:"boolean"}
def show_pdb(pdb_file, flag=0, show_sidechains=False, show_mainchains=False, color="lDDT"):
model_name = f"Flag_{flag}"
view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)
view.addModel(open(pdb_file,'r').read(),'pdb')
if color == "lDDT":
view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':50,'max':90}}})
elif color == "rainbow":
view.setStyle({'cartoon': {'color':'spectrum'}})
elif color == "chain":
chains = len(queries[0][1]) + 1 if is_complex else 1
for n,chain,color in zip(range(chains),list("ABCDEFGH"),
["lime","cyan","magenta","yellow","salmon","white","blue","orange"]):
view.setStyle({'chain':chain},{'cartoon': {'color':color}})
if show_sidechains:
BB = ['C','O','N']
view.addStyle({'and':[{'resn':["GLY","PRO"],'invert':True},{'atom':BB,'invert':True}]},
{'stick':{'colorscheme':f"WhiteCarbon",'radius':0.3}})
view.addStyle({'and':[{'resn':"GLY"},{'atom':'CA'}]},
{'sphere':{'colorscheme':f"WhiteCarbon",'radius':0.3}})
view.addStyle({'and':[{'resn':"PRO"},{'atom':['C','O'],'invert':True}]},
{'stick':{'colorscheme':f"WhiteCarbon",'radius':0.3}})
if show_mainchains:
BB = ['C','O','N','CA']
view.addStyle({'atom':BB},{'stick':{'colorscheme':f"WhiteCarbon",'radius':0.3}})
view.zoomTo()
if color == "lDDT":
plot_plddt_legend().show()
return view
def get_avg_Bfac (file='./output_v3/[0.0 0.5 0.0 0.0 0.0 0.0 0.0 0.0].pdb'):
p = PDBParser()
avg_B=0
bfac_list=[]
structure = p.get_structure("X", file)
for PDBmodel in structure:
for chain in PDBmodel:
for residue in chain:
for atom in residue:
Bfac=atom.get_bfactor()
bfac_list.append(Bfac)
avg_B=avg_B+Bfac
avg_B=avg_B/len (bfac_list)
print (f"For {file}, average B-factor={avg_B}")
plt.plot (bfac_list, label='lDDT')
plt.xlabel ('Atom #' )
plt.ylabel ('iDDT')
plt.legend()
plt.show()
return avg_B, bfac_list
def sample_sequence_normalized_Bfac (seccs=[0.3, 0.3, 0.1, 0., 0., 0., 0., 0. ]):
sample_numbers=torch.tensor([seccs])
sample_numbers=torch.nn.functional.normalize (sample_numbers, dim=1)
sample_numbers=sample_numbers/torch.sum(sample_numbers)
PDB=sample_sequence (model,
X=sample_numbers,
flag=0,cond_scales=1, foldproteins=True,
)
avg,_ = get_avg_Bfac (file=PDB[0])
return PDB, avg
# ======================================================
# blocks for Model A
# ======================================================
def train_loop_Old_FromModelA (
model,
train_loader,
test_loader,
#
optimizer=None,
print_every=1,
epochs= 300,
start_ep=0,
start_step=0,
train_unet_number=1,
print_loss_every_steps=1000,
#
trainer=None,
plot_unscaled=False,
max_batch_size=4,
save_model=False,
cond_scales=[1.0], #list of cond scales
num_samples=2, #how many samples produced every time tested.....
foldproteins=False,
# ++
cond_image=False, # not use cond_images... for model A
cond_text=True, # use condi_text... for model A
# +
device=None,
loss_list=[],
epoch_list=[],
train_hist_file=None,
train_hist_file_full=None,
prefix=None, # not used in this function
Xnormfac=None,
ynormfac=1.,
tokenizer_X=None,
tokenizer_y=None,
test_condition_list=[],
max_length_Y=1,
max_text_len_X=1,
CKeys=None,
sample_steps=1,
sample_dir=None,
save_every_epoch=1,
save_point_info_file=None,
store_dir=None,
):
# #+
# Xnormfac=Xnormfac.to(model.device)
if not exists (trainer):
if not exists (optimizer):
print ("ERROR: If trainer not used, need to provide optimizer.")
if exists (trainer):
print ("Trainer provided... will be used")
# --------------------------------
# steps=start_step
# ++++++++++++++++++++++++++++++++
steps=start_step
added_steps=0
loss_total=0
for e in range(1, epochs+1):
# start = time.time()
torch.cuda.empty_cache()
print ("######################################################################################")
start = time.time()
print ("NOW: Training epoch: ", e+start_ep)
# TRAINING
train_epoch_loss = 0
model.train()
print ("Loop over ", len(train_loader), " batches (print . every ", print_every, " steps)")
for item in train_loader:
# ++
steps += 1
added_steps += 1
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
if exists (trainer):
if cond_image==False:
# ========================================
# Model A: condition via text
# ========================================
# this block depends on the model:forward
loss = trainer(
# # --------------------------------
# X_train_batch,
# y_train_batch.unsqueeze(1) ,
# ++++++++++++++++++++++++++++++++
y_train_batch.unsqueeze(1) ,
x=X_train_batch,
#
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
if cond_image==True:
# ========================================
# Model B: condition via image/sequence
# ========================================
# added for future: Train_loop B
loss = trainer(
y_train_batch.unsqueeze(1) ,
x=None,
cond_images=X_train_batch.unsqueeze(1),
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
# pass
#
trainer.update(unet_number = train_unet_number)
else:
optimizer.zero_grad()
if cond_image==False:
# this block depends on the model:forward
loss=model (
# # --------------------------------
# X_train_batch,
# y_train_batch.unsqueeze(1) ,
# ++++++++++++++++++++++++++++++++
y_train_batch.unsqueeze(1) ,
x=X_train_batch,
#
unet_number=train_unet_number
)
if cond_image==True:
# added for future: Train_loop B
loss=model (
y_train_batch.unsqueeze(1) ,
x=None,
cond_images=X_train_batch.unsqueeze(1),
unet_number=train_unet_number
)
#
loss.backward( )
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
loss_total=loss_total+loss.item()
# +
train_epoch_loss=train_epoch_loss+loss.item()
if steps % print_every == 0:
# for progress bar
print(".", end="")
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# record loss block
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# if steps>0:
if added_steps>0:
if steps % print_loss_every_steps == 0:
# + for debug
if CKeys['Debug_TrainerPack']==2:
print("Here is step: ", steps)
norm_loss=loss_total/print_loss_every_steps
print (f"\nTOTAL LOSS at epoch={e+start_ep}, step={steps}: {norm_loss}")
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# add a line to the hist file
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss)+'\n'
with open(train_hist_file,'a') as f:
f.write(add_line)
loss_list.append (norm_loss)
loss_total=0
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
epoch_list.append(e+start_ep)
fig = plt.figure()
plt.plot (epoch_list, loss_list, label='Loss')
plt.legend()
# outname = prefix+ f"loss_{e}_{steps}.jpg"
outname = sample_dir+ f"loss_{e+start_ep}_{steps}.jpg"
#
# the order, save then show, matters
if CKeys['SlientRun']==1:
plt.savefig(outname, dpi=200)
else:
plt.show()
plt.close(fig)
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# sample in test set block
# set sample_steps < 0 to switch off this block
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# if steps>0:
if added_steps>0:
if steps % sample_steps == 0 and sample_steps > 0:
# + for debug
if CKeys['Debug_TrainerPack']==2:
print("Here is steps: ", steps)
if plot_unscaled:
# test before scaling...
plt.plot (
y_train_batch.unsqueeze(1)[0,0,:].cpu().detach().numpy(),
label= 'Unscaled GT'
)
plt.legend()
plt.show()
#rescale GT to properly plot
GT=y_train_batch.cpu().detach()
GT=resize_image_to(
GT.unsqueeze(1),
model.imagen.image_sizes[train_unet_number-1],
)
####
print ("I. SAMPLING IN TEST SET: ")
####
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
sample_loop_omegafold_ModelA (
model,
test_loader,
cond_scales=cond_scales,
num_samples=num_samples, #how many samples produced every time tested.....
timesteps=None,
flag=e+start_ep, # steps,
foldproteins=foldproteins,
# add condi_key
cond_image=cond_image, # Not used for now
cond_text=cond_text, # Not used for now
skip_steps=0,
#
max_text_len=max_text_len_X,
max_length=max_length_Y,
# ++++++++++++++++++++
train_unet_number=train_unet_number,
ynormfac=ynormfac,
prefix=prefix, #
tokenizer_y=tokenizer_y,
Xnormfac_CondiText=Xnormfac,
tokenizer_X_CondiText=tokenizer_X,
# ++
CKeys=CKeys,
sample_dir=sample_dir,
steps=steps,
e=e+start_ep,
IF_showfig= CKeys['SlientRun']!=1 ,
)
print ("II. SAMPLING FOR DE NOVO:")
sample_sequence_omegafold_ModelA (
# # ----------------------------------------------
# model,
# X=[[0, 0.7, 0.07, 0.1, 0.01, 0.02, 0.01, 0.11]],
# foldproteins=foldproteins,
# flag=steps,cond_scales=1.,
# ++++++++++++++++++++++++++++++++++++++++++++++
model,
X=test_condition_list, # [[0.92, 0., 0.04, 0.04, 0., 0., 0., 0., ]], # from text conditioning X
flag=e+start_ep, # steps, # 0,
cond_scales=cond_scales, # 1.,
foldproteins=True, # False,
X_string=None, # from text conditioning X_string
x_data=None, # from image conditioning x_data
skip_steps=0,
inpaint_images=None, # in formation Y data
inpaint_masks = None,
inpaint_resample_times = None,
init_images = None,
num_cycle=16, # for omegafolding
calc_error=False, # for check on folded results, not used for every case
# ++++++++++++++++++++++++++
# tokenizers
tokenizer_X_forImageCondi=None, # for x_data
Xnormfac_forImageCondi=1.,
tokenizer_X_forTextCondi=None, # for X if NEEDED only
Xnormfac_forTextCondi=1.,
tokenizer_y=tokenizer_y, # None, # for output Y
ynormfac=ynormfac,
# length
train_unet_number=1,
max_length_Y=max_length_Y, # for Y, X_forImageCondi
max_text_len=max_text_len_X, # for X_forTextCondi
# other info
steps=steps, # None,
e=e, # None,
sample_dir=sample_dir, # None,
prefix=prefix, # None,
IF_showfig= CKeys['SlientRun']!=1, # True,
CKeys=CKeys,
# TBA to Model B
normalize_X_cond_to_one=False,
)
# sample_sequence (model,
# X=[[0., 0.0, 0.0, 0.0, 0., 0., 0., 0., ]],foldproteins=foldproteins,
# flag=steps,cond_scales=1.,
# )
# summerize loss over every epoch:
norm_loss_over_e = train_epoch_loss/len(train_loader)
print("\nnorm_loss over 1 epoch: ", norm_loss_over_e)
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
# write this into "train_hist_file_full"
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss_over_e)+'\n'
with open(train_hist_file_full,'a') as f:
f.write(add_line)
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# save model every this epoches
if save_model and (e+start_ep) % save_every_epoch==0 and e>1:
# fname=f"{prefix}trainer_save-model-epoch_{e+start_ep}.pt"
fname=f"{store_dir}trainer_save-model-epoch_{e+start_ep}.pt"
trainer.save(fname)
print (f"Model saved: ", fname)
# fname=f"{prefix}statedict_save-model-epoch_{e+start_ep}.pt"
fname=f"{store_dir}statedict_save-model-epoch_{e+start_ep}.pt"
torch.save(model.state_dict(), fname)
print (f"Statedict model saved: ", fname)
# add a saving point file
top_line='epoch,steps,norm_loss'+'\n'
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss)+'\n'
with open(save_point_info_file, "w") as f:
f.write(top_line)
f.write(add_line)
# if steps>0:
# if save_model and steps % print_loss_every_steps==0:
# fname=f"{prefix}trainer_save-model-epoch_{e}.pt"
# trainer.save(fname)
# fname=f"{prefix}statedict_save-model-epoch_{e}.pt"
# torch.save(model.state_dict(), fname)
# print (f"Model saved: ")
# steps=steps+1
print (f"\n\n-------------------\nTime for epoch {e+start_ep}={(time.time()-start)/60}\n-------------------")
def train_loop_ForModelA_II (
model,
train_loader,
test_loader,
#
optimizer=None,
print_every=1,
epochs= 300,
start_ep=0,
start_step=0,
train_unet_number=1,
print_loss_every_steps=1000,
#
trainer=None,
plot_unscaled=False,
max_batch_size=4,
save_model=False,
cond_scales=[1.0], #list of cond scales
num_samples=2, #how many samples produced every time tested.....
foldproteins=False,
# ++
cond_image=False, # not use cond_images... for model A
cond_text=True, # use condi_text... for model A
# +
device=None,
loss_list=[],
epoch_list=[],
train_hist_file=None,
train_hist_file_full=None,
prefix=None, # not used in this function
Xnormfac=None,
ynormfac=1.,
tokenizer_X=None,
tokenizer_y=None,
test_condition_list=[],
max_length_Y=1,
max_text_len_X=1,
CKeys=None,
sample_steps=1,
sample_dir=None,
save_every_epoch=1,
save_point_info_file=None,
store_dir=None,
# ++ for pLM
pLM_Model_Name=None,
image_channels=None,
print_error=False, # not defined for Problem6 # True,
):
# #+
# Xnormfac=Xnormfac.to(model.device)
if not exists (trainer):
if not exists (optimizer):
print ("ERROR: If trainer not used, need to provide optimizer.")
if exists (trainer):
print ("Trainer provided... will be used")
# --------------------------------
# steps=start_step
# ++++++++++++++++++++++++++++++++
steps=start_step
added_steps=0
loss_total=0
# ++ for pLM
if pLM_Model_Name=='None':
pLM_Model=None
elif pLM_Model_Name=='esm2_t33_650M_UR50D':
# dim: 1280
esm_layer=33
pLM_Model, esm_alphabet = esm.pretrained.esm2_t33_650M_UR50D()
len_toks=len(esm_alphabet.all_toks)
pLM_Model.eval()
pLM_Model. to(device)
elif pLM_Model_Name=='esm2_t36_3B_UR50D':
# dim: 2560
esm_layer=36
pLM_Model, esm_alphabet = esm.pretrained.esm2_t36_3B_UR50D()
len_toks=len(esm_alphabet.all_toks)
pLM_Model.eval()
pLM_Model. to(device)
elif pLM_Model_Name=='esm2_t30_150M_UR50D':
# dim: 640
esm_layer=30
pLM_Model, esm_alphabet = esm.pretrained.esm2_t30_150M_UR50D()
len_toks=len(esm_alphabet.all_toks)
pLM_Model.eval()
pLM_Model. to(device)
elif pLM_Model_Name=='esm2_t12_35M_UR50D':
# dim: 480
esm_layer=12
pLM_Model, esm_alphabet = esm.pretrained.esm2_t12_35M_UR50D()
len_toks=len(esm_alphabet.all_toks)
pLM_Model.eval()
pLM_Model. to(device)
else:
print("pLM model is missing...")
for e in range(1, epochs+1):
# start = time.time()
torch.cuda.empty_cache()
print ("######################################################################################")
start = time.time()
print ("NOW: Training epoch: ", e+start_ep)
# TRAINING
train_epoch_loss = 0
model.train()
print ("Loop over ", len(train_loader), " batches (print . every ", print_every, " steps)")
for item in train_loader:
# ++
steps += 1
added_steps += 1
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
# project y_ into embedding space
if CKeys["Debug_TrainerPack"]==1:
print("Initial unload the dataloader items: ...")
print(X_train_batch.shape)
print(y_train_batch.shape)
# ++
# project the AA seq into embedding space
# for output, it is shared between ModelA and ModelB
# # --
# if pLM_Model_Name=='None':
# # just use the encoded sequence
# y_train_batch_in = y_train_batch.unsqueeze(1)
# # pass
# elif pLM_Model_Name=='esm2_t33_650M_UR50D':
# with torch.no_grad():
# results = pLM_Model(
# y_train_batch,
# repr_layers=[33],
# return_contacts=False,
# )
# y_train_batch_in = results["representations"][33]
# y_train_batch_in = rearrange(
# y_train_batch_in,
# 'b l c -> b c l'
# )
# else:
# print(f"Required pLM name is not defined!!")
# ++
if pLM_Model_Name=='None':
# just use the encoded sequence
y_train_batch_in = y_train_batch.unsqueeze(1)
# pass
else: # for ESM models # pLM_Model_Name=='esm2_t33_650M_UR50D':
with torch.no_grad():
results = pLM_Model(
y_train_batch,
repr_layers=[esm_layer],
return_contacts=False,
)
y_train_batch_in = results["representations"][esm_layer]
y_train_batch_in = rearrange(
y_train_batch_in,
'b l c -> b c l'
)
#
# For input part, this block is different for ModelA and ModelB
if cond_image==False:
# model A: X: text_condi, not affected by pLM
X_train_batch_in = X_train_batch
else:
# model B: X: cond_img, will be affected by pLM
X_train_batch_in = X_train_batch.unsqueeze(1).repeat(1,image_channels,1)
#
# + for debug
if CKeys["Debug_TrainerPack"]==1:
print("After pLM model, the shape of X and y for training:")
print("X_train_batch_in.dim: ", X_train_batch_in.shape)
print("y_train_batch_in.dim: ", y_train_batch_in.shape)
if exists (trainer):
if cond_image==False:
# ========================================
# Model A: condition via text
# ========================================
# this block depends on the model:forward
loss = trainer(
# # --------------------------------
# X_train_batch,
# y_train_batch.unsqueeze(1) ,
# # ++++++++++++++++++++++++++++++++
# y_train_batch.unsqueeze(1) ,
# x=X_train_batch,
# ++ pLM
y_train_batch_in,
x=X_train_batch_in,
#
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
if cond_image==True:
# ========================================
# Model B: condition via image/sequence
# ========================================
# # --
# # added for future: Train_loop B
# loss = trainer(
# y_train_batch.unsqueeze(1) ,
# x=None,
# cond_images=X_train_batch.unsqueeze(1),
# unet_number=train_unet_number,
# max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
# )
# ++ from pLM+ModelB
loss = trainer(
y_train_batch_in, # true image
x=None, # tokenized text
cond_images=X_train_batch_in, # cond_image
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
# pass
#
trainer.update(unet_number = train_unet_number)
else:
optimizer.zero_grad()
if cond_image==False:
# this block depends on the model:forward
loss=model (
# # --------------------------------
# X_train_batch,
# y_train_batch.unsqueeze(1) ,
# # ++++++++++++++++++++++++++++++++
# y_train_batch.unsqueeze(1) ,
# x=X_train_batch,
# ++ pLM
y_train_batch_in,
x=X_train_batch_in,
#
unet_number=train_unet_number
)
if cond_image==True:
# added for future: Train_loop B
# # --
# loss=model (
# y_train_batch.unsqueeze(1) ,
# x=None,
# cond_images=X_train_batch.unsqueeze(1),
# unet_number=train_unet_number
# )
# ++ from pLM
loss=model (
y_train_batch_in ,
x=None,
cond_images=X_train_batch_in,
unet_number=train_unet_number
)
#
loss.backward( )
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
loss_total=loss_total+loss.item()
# +
train_epoch_loss=train_epoch_loss+loss.item()
if steps % print_every == 0:
# for progress bar
print(".", end="")
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# record loss block
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# if steps>0:
if added_steps>0:
if steps % print_loss_every_steps == 0:
# + for debug
if CKeys['Debug_TrainerPack']==2:
print("Here is step: ", steps)
norm_loss=loss_total/print_loss_every_steps
print (f"\nTOTAL LOSS at epoch={e+start_ep}, step={steps}: {norm_loss}")
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# add a line to the hist file
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss)+'\n'
with open(train_hist_file,'a') as f:
f.write(add_line)
loss_list.append (norm_loss)
loss_total=0
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
epoch_list.append(e+start_ep)
fig = plt.figure()
plt.plot (epoch_list, loss_list, label='Loss')
plt.legend()
# outname = prefix+ f"loss_{e}_{steps}.jpg"
outname = sample_dir+ f"loss_{e+start_ep}_{steps}.jpg"
#
# the order, save then show, matters
if CKeys['SlientRun']==1:
plt.savefig(outname, dpi=200)
else:
plt.show()
plt.close(fig)
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# sample in test set block
# set sample_steps < 0 to switch off this block
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# if steps>0:
if added_steps>0:
if steps % sample_steps == 0 and sample_steps > 0:
# + for debug
if CKeys['Debug_TrainerPack']==2:
print("Here is steps: ", steps)
if plot_unscaled:
# test before scaling...
plt.plot (
y_train_batch.unsqueeze(1)[0,0,:].cpu().detach().numpy(),
label= 'Unscaled GT'
)
plt.legend()
plt.show()
# # -- look like not used
# #rescale GT to properly plot
# GT=y_train_batch.cpu().detach()
# GT=resize_image_to(
# GT.unsqueeze(1),
# model.imagen.image_sizes[train_unet_number-1],
# )
####
print (">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ")
print ("I. SAMPLING IN TEST SET: ")
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
####
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
sample_loop_omegafold_pLM_ModelA (
model,
test_loader,
cond_scales=cond_scales,
num_samples=num_samples, #how many samples produced every time tested.....
timesteps=None,
flag=e+start_ep, # steps,
foldproteins=foldproteins,
# add condi_key
cond_image=cond_image, # Not used for now
cond_text=cond_text, # Not used for now
skip_steps=0,
#
max_text_len=max_text_len_X,
max_length=max_length_Y,
# ++++++++++++++++++++
train_unet_number=train_unet_number,
ynormfac=ynormfac,
prefix=prefix, #
tokenizer_y=tokenizer_y,
Xnormfac_CondiText=Xnormfac,
tokenizer_X_CondiText=tokenizer_X,
# ++
CKeys=CKeys,
sample_dir=sample_dir,
steps=steps,
e=e+start_ep,
IF_showfig= CKeys['SlientRun']!=1 ,
# ++ for pLM
pLM_Model=pLM_Model,
pLM_Model_Name=pLM_Model_Name,
image_channels=image_channels,
pLM_alphabet=esm_alphabet,
)
print (">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> ")
print ("II. SAMPLING FOR DE NOVO:")
print ("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< ")
DeNovoSam_pdbs, fasta_file_list=\
sample_sequence_omegafold_pLM_ModelA (
# # ----------------------------------------------
# model,
# X=[[0, 0.7, 0.07, 0.1, 0.01, 0.02, 0.01, 0.11]],
# foldproteins=foldproteins,
# flag=steps,cond_scales=1.,
# ++++++++++++++++++++++++++++++++++++++++++++++
model,
X=test_condition_list, # [[0.92, 0., 0.04, 0.04, 0., 0., 0., 0., ]], # from text conditioning X
flag=e+start_ep, # steps, # 0,
cond_scales=cond_scales, # 1.,
foldproteins=True, # False,
X_string=None, # from text conditioning X_string
x_data=None, # from image conditioning x_data
skip_steps=0,
inpaint_images=None, # in formation Y data
inpaint_masks = None,
inpaint_resample_times = None,
init_images = None,
num_cycle=16, # for omegafolding
calc_error=False, # for check on folded results, not used for every case
# ++++++++++++++++++++++++++
# tokenizers
tokenizer_X_forImageCondi=None, # for x_data
Xnormfac_forImageCondi=1.,
tokenizer_X_forTextCondi=None, # for X if NEEDED only
Xnormfac_forTextCondi=1.,
tokenizer_y=tokenizer_y, # None, # for output Y
ynormfac=ynormfac,
# length
train_unet_number=1,
max_length_Y=max_length_Y, # for Y, X_forImageCondi
max_text_len=max_text_len_X, # for X_forTextCondi
# other info
steps=steps, # None,
e=e, # None,
sample_dir=sample_dir, # None,
prefix=prefix, # None,
IF_showfig= CKeys['SlientRun']!=1, # True,
CKeys=CKeys,
# TBA to Model B
normalize_X_cond_to_one=False,
# ++ for pLM
pLM_Model=pLM_Model,
pLM_Model_Name=pLM_Model_Name,
image_channels=image_channels,
pLM_alphabet=esm_alphabet,
)
# sample_sequence (model,
# X=[[0., 0.0, 0.0, 0.0, 0., 0., 0., 0., ]],foldproteins=foldproteins,
# flag=steps,cond_scales=1.,
# )
# summerize loss over every epoch:
norm_loss_over_e = train_epoch_loss/len(train_loader)
print("\nnorm_loss over 1 epoch: ", norm_loss_over_e)
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
# write this into "train_hist_file_full"
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss_over_e)+'\n'
with open(train_hist_file_full,'a') as f:
f.write(add_line)
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# save model every this epoches
if save_model and (e+start_ep) % save_every_epoch==0 and e>1:
# fname=f"{prefix}trainer_save-model-epoch_{e+start_ep}.pt"
fname=f"{store_dir}trainer_save-model-epoch_{e+start_ep}.pt"
trainer.save(fname)
print (f"Model saved: ", fname)
# fname=f"{prefix}statedict_save-model-epoch_{e+start_ep}.pt"
fname=f"{store_dir}statedict_save-model-epoch_{e+start_ep}.pt"
torch.save(model.state_dict(), fname)
print (f"Statedict model saved: ", fname)
# add a saving point file
top_line='epoch,steps,norm_loss'+'\n'
add_line = str(e+start_ep)+','+str(steps)+','+str(norm_loss)+'\n'
with open(save_point_info_file, "w") as f:
f.write(top_line)
f.write(add_line)
# if steps>0:
# if save_model and steps % print_loss_every_steps==0:
# fname=f"{prefix}trainer_save-model-epoch_{e}.pt"
# trainer.save(fname)
# fname=f"{prefix}statedict_save-model-epoch_{e}.pt"
# torch.save(model.state_dict(), fname)
# print (f"Model saved: ")
# steps=steps+1
print (f"\n\n-------------------\nTime for epoch {e+start_ep}={(time.time()-start)/60}\n-------------------")
# from original, not used any more
def train_loop_Model_A (
model,
train_loader,
test_loader,
optimizer=None,
print_every=10,
epochs= 300,
start_ep=0,
start_step=0,
train_unet_number=1,
print_loss=1000,
trainer=None,
plot_unscaled=False,
max_batch_size=4,
save_model=False,
cond_scales=[1.0], #list of cond scales
num_samples=2, #how many samples produced every time tested.....
foldproteins=False,
):
if not exists (trainer):
if not exists (optimizer):
print ("ERROR: If trainer not used, need to provide optimizer.")
if exists (trainer):
print ("Trainer provided... will be used")
steps=start_step
loss_total=0
for e in range(1, epochs+1):
start = time.time()
torch.cuda.empty_cache()
print ("######################################################################################")
start = time.time()
print ("NOW: Training epoch: ", e+start_ep)
# TRAINING
train_epoch_loss = 0
model.train()
print ("Loop over ", len(train_loader), " batches (print . every ", print_every, " steps)")
for item in train_loader:
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
if exists (trainer):
loss = trainer(
X_train_batch, y_train_batch.unsqueeze(1) ,
unet_number=train_unet_number,
max_batch_size = max_batch_size, # auto divide the batch of 64 up into batch size of 4 and accumulate gradients, so it all fits in memory
)
trainer.update(unet_number = train_unet_number)
else:
optimizer.zero_grad()
loss=model ( X_train_batch, y_train_batch.unsqueeze(1) ,unet_number=train_unet_number)
loss.backward( )
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
loss_total=loss_total+loss.item()
if steps % print_every == 0:
print(".", end="")
if steps>0:
if steps % print_loss == 0:
if plot_unscaled:
plt.plot (y_train_batch.unsqueeze(1)[0,0,:].cpu().detach().numpy(),label= 'Unscaled GT')
plt.legend()
plt.show()
#rescale GT to properly plot
GT=y_train_batch.cpu().detach()
GT=resize_image_to(
GT.unsqueeze(1),
model.imagen.image_sizes[train_unet_number-1],
)
norm_loss=loss_total/print_loss
print (f"\nTOTAL LOSS at epoch={e}, step={steps}: {norm_loss}")
loss_list.append (norm_loss)
loss_total=0
plt.plot (loss_list, label='Loss')
plt.legend()
outname = prefix+ f"loss_{e}_{steps}.jpg"
plt.savefig(outname, dpi=200)
plt.show()
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
sample_loop (model,
test_loader,
cond_scales=cond_scales,
num_samples=1, #how many samples produced every time tested.....
timesteps=64,
flag=steps,foldproteins=foldproteins,
)
print ("SAMPLING FOR DE NOVO:")
sample_sequence (model,
X=[[0, 0.7, 0.07, 0.1, 0.01, 0.02, 0.01, 0.11]],foldproteins=foldproteins,
flag=steps,cond_scales=1.,
)
sample_sequence (model,
X=[[0., 0.0, 0.0, 0.0, 0., 0., 0., 0., ]],foldproteins=foldproteins,
flag=steps,cond_scales=1.,
)
if steps>0:
if save_model and steps % print_loss==0:
fname=f"{prefix}trainer_save-model-epoch_{e}.pt"
trainer.save(fname)
fname=f"{prefix}statedict_save-model-epoch_{e}.pt"
torch.save(model.state_dict(), fname)
print (f"Model saved: ")
steps=steps+1
print (f"\n\n-------------------\nTime for epoch {e}={(time.time()-start)/60}\n-------------------")
# +++
def sample_sequence_omegafold_ModelA (
model,
X=[[0.92, 0., 0.04, 0.04, 0., 0., 0., 0., ]], # from text conditioning X
flag=0,
cond_scales=1.,
foldproteins=False,
X_string=None, # from text conditioning X_string
x_data=None, # from image conditioning x_data
skip_steps=0,
inpaint_images=None, # in formation Y data
inpaint_masks = None,
inpaint_resample_times = None,
init_images = None,
num_cycle=16, # for omegafolding
calc_error=False, # for check on folded results, not used for every case
# ++++++++++++++++++++++++++
# tokenizers
tokenizer_X_forImageCondi=None, # for x_data
Xnormfac_forImageCondi=1.,
tokenizer_X_forTextCondi=None, # for X if NEEDED only
Xnormfac_forTextCondi=1.,
tokenizer_y=None, # for output Y
ynormfac=1,
# length
train_unet_number=1,
max_length_Y=1, # for Y, X_forImageCondi
max_text_len=1, # for X_forTextCondi
# other info
steps=None,
e=None,
sample_dir=None,
prefix=None,
IF_showfig=True,
CKeys=None,
# TBA to Model B
normalize_X_cond_to_one=False,
):
# -----------
# steps=0
# e=flag
# --
# print (f"Producing {len(X)} samples...")
# ++
if X!=None:
print (f"Producing {len(X)} samples...from text conditioning X...")
lenn_val=len(X)
if X_string!=None:
lenn_val=len(X_string)
print (f"Producing {len(X_string)} samples...from text conditioning X_String (from string)...")
if x_data!=None:
print (f"Producing {len(x_data)} samples...from image conditingig x_data ...")
lenn_val=len(x_data)
# print (x_data)
print ('Device: ', model.device)
for iisample in range (lenn_val):
print(f"Working on {iisample}")
X_cond=None
if X_string==None and X!=None: # for X channel
X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
if X_string!=None: # from raw text, ie., X_string: need tokenizer_X and Xnormfac
# -
# X = tokenizer_X.texts_to_sequences(X_string[iisample])
# X= sequence.pad_sequences(X, maxlen=max_length, padding='post', truncating='post')
# X=np.array(X)
# X_cond=torch.from_numpy(X).float()/Xnormfac
# +
XX = tokenizer_X_forTextCondi.texts_to_sequences(X_string[iisample])
XX = sequence.pad_sequences(XX, maxlen=max_text_len, padding='post', truncating='post')
XX = np.array(XX)
X_cond = torch.from_numpy(XX).float()/Xnormfac_forTextCondi
print ('Tokenized and processed: ', X_cond)
if X_cond!=None:
if normalize_X_cond_to_one: # used when there is constrain on X_cond.sum()
X_cond=X_cond/X_cond.sum()
print ("Text conditoning used: ", X_cond, "...sum: ", X_cond.sum(), "cond scale: ", cond_scales)
else:
print ("Text conditioning used: None")
# for now, assume image_condi and text_condi can be used at the same time
if tokenizer_X_forImageCondi==None:
# ===========================================================
# condi_image/seq needs no tokenization, like numbers: force_path
# only normalization needed
# Based on ModelB:Force_Path
if x_data!=None:
x_data_tokenized=torch.from_numpy(x_data[iisample]/Xnormfac_forImageCondi)
x_data_tokenized=x_data_tokenized.to(torch.float)
# + for debug:
if CKeys['Debug_TrainerPack']==1:
print("x_data_tokenized dim: ", x_data_tokenized.shape)
print("x_data_tokenized dtype: ", x_data_tokenized.dtype)
print("test: ", x_data_tokenized!=None)
else:
x_data_tokenized=None
# + for debug:
if CKeys['Debug_TrainerPack']==1:
print("x_data_tokenized and x_data: None")
# model.sample:full arguments
# self,
# x=None,
# stop_at_unet_number=1,
# cond_scale=7.5,
# # ++
# x_data=None, # image_condi data
# skip_steps=None,
# inpaint_images = None,
# inpaint_masks = None,
# inpaint_resample_times = 5,
# init_images = None,
# x_data_tokenized=None,
# tokenizer_X=None,
# Xnormfac=1.,
# # -+
# device=None,
# max_length=1., # for XandY data, in image/sequence format; NOT for text condition
# max_text_len=1., # for X data, in text format
#
result=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales ,
x_data=None,
# ++
x_data_tokenized=x_data_tokenized,
#
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,
device=model.device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X_forImageCondi, # tokenizer_X,
Xnormfac=Xnormfac_forImageCondi, # Xnormfac,
# ynormfac=ynormfac,
max_length=max_length_Y, # for ImageCondi, max_length,
max_text_len=max_text_len,
)
else:
#
result=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales ,
x_data=x_data[iisample],
# ++
x_data_tokenized=None,
#
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,
device=model.device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X_forImageCondi, # tokenizer_X,
Xnormfac=Xnormfac_forImageCondi, # Xnormfac,
# ynormfac=ynormfac,
max_length=max_length_Y, # max_length,
max_text_len=max_text_len,
)
# # ------------------------------------------
# result=model.sample (
# X_cond,
# stop_at_unet_number=train_unet_number,
# cond_scale=cond_scales
# )
result=torch.round(result*ynormfac)
# + for debug
print("result.dim: ", result.shape)
fig=plt.figure()
plt.plot (
result[0,0,:].cpu().detach().numpy(),
label= f'Predicted'
)
#plt.plot (GT[samples,0,:]*ynormfac,label= f'GT {0}')
plt.legend()
outname = sample_dir+ f"sampled_from_X_{iisample}_condscale-{str (cond_scales)}_{e}_{steps}.jpg"
#plt.title (f"Sample {samples}, cond scale={str (cond_scales[iisample])}")
if IF_showfig==1:
plt.show ()
else:
plt.savefig(outname, dpi=200)
plt.close()
# # ----------------------------------------
# plt.plot (result[0,0,:].cpu().detach().numpy(),label= f'Predicted')
# plt.legend()
# outname = prefix+ f"sampld_from_X_{flag}_condscale-{str (cond_scales)}_{e}_{steps}.jpg"
# plt.savefig(outname, dpi=200)
# plt.show ()
to_rev=result[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
print("to_rev.dim: ", to_rev.shape)
y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
# + from Model B
### reverse second structure input....
pdb_list=[]
if X_cond != None:
# there is condi_text
if X_string!=None:
X_cond=torch.round(X_cond*Xnormfac_forTextCondi)
to_rev=X_cond[:,:]
to_rev=to_rev.long().cpu().detach().numpy()
print ("to_rev.dim: ", to_rev.shape)
# --
# X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
# ++
X_text_reversed=tokenizer_X_forTextCondi.sequences_to_texts (to_rev)
for iii in range (len(y_text_reversed)):
X_text_reversed[iii]=X_text_reversed[iii].upper().strip().replace(" ", "")
if X_string==None:
# reverse this: X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
X_text_reversed=X_cond
else:
X_text_reversed=None
if x_data !=None: # there is condi_image
x_data_reversed=x_data #is already in sequence fromat..
else:
x_data_reversed=None
# summary
# print (f"For {X_text_reversed} or {X[iisample]} on Text_Condi,\n and {x_data_reversed} on Image_Condi,\n predicted sequence: ", y_data_reversed)
print (f"For {X_text_reversed} or {X[iisample]} on Text_Condi,\n and {x_data_reversed} on Image_Condi,")
print (f"predicted sequence full: {y_data_reversed}")
# add just for incase check
print (f"predicted sequence: {y_data_reversed[0]}")
# + for debug
print("================================================")
print("foldproteins: ", foldproteins)
if not foldproteins:
pdb_file=None
else:
# if foldproteins:
if X_cond != None:
pass
tempname='temp'
pdb_file=foldandsavePDB (
sequence=y_data_reversed[0],
filename_out=tempname,
num_cycle=num_cycle,
flag=flag,
# +++++++++++++++++++
# prefix=prefix,
prefix=sample_dir,
)
#
out_nam=iisample
#
out_nam_fasta=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.fasta'
write_fasta (y_data_reversed[0], out_nam_fasta)
#
out_nam=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.pdb'
shutil.copy (pdb_file, out_nam) #source, dest
pdb_file=out_nam
#
print (f"Properly named PDB file produced: {pdb_file}")
if IF_showfig==1:
#flag=1000
view=show_pdb(
pdb_file=pdb_file,
flag=flag,
show_sidechains=show_sidechains,
show_mainchains=show_mainchains,
color=color
)
view.show()
if calc_error:
# only work for ModelA:SecStr
if CKeys['Problem_ID']==7:
get_Model_A_error (pdb_file, X[iisample], plotit=True)
else:
print ("Error calculation on the predicted results is not applicable")
pdb_list.append(pdb_file)
return pdb_list
# xbc=X_cond[iisample,:].cpu().detach().numpy()
# out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.2f" % xbc})+f'_{flag}_{steps}'
# tempname='temp'
# pdb_file=foldandsavePDB (sequence=y_data_reversed[0],
# filename_out=tempname,
# num_cycle=16, flag=flag)
# out_nam_fasta=f'{prefix}{out_nam}.fasta'
# out_nam=f'{prefix}{out_nam}.pdb'
# write_fasta (y_data_reversed[0], out_nam_fasta)
# shutil.copy (pdb_file, out_nam) #source, dest
# pdb_file=out_nam
# print (f"Properly named PDB file produced: {pdb_file}")
# view=show_pdb(pdb_file=pdb_file, flag=flag,
# show_sidechains=show_sidechains, show_mainchains=show_mainchains, color=color)
# view.show()
# if calc_error:
# get_Model_A_error (pdb_file, X[iisample], plotit=True)
# return pdb_file
# +++
def sample_sequence_omegafold_pLM_ModelA (
model,
X=[[0.92, 0., 0.04, 0.04, 0., 0., 0., 0., ]], # from text conditioning X
flag=0,
cond_scales=1.,
foldproteins=False,
X_string=None, # from text conditioning X_string
x_data=None, # from image conditioning x_data
skip_steps=0,
inpaint_images=None, # in formation Y data
inpaint_masks = None,
inpaint_resample_times = None,
init_images = None,
num_cycle=16, # for omegafolding
calc_error=False, # for check on folded results, not used for every case
# ++++++++++++++++++++++++++
# tokenizers
tokenizer_X_forImageCondi=None, # for x_data
Xnormfac_forImageCondi=1.,
tokenizer_X_forTextCondi=None, # for X if NEEDED only
Xnormfac_forTextCondi=1.,
tokenizer_y=None, # for output Y
ynormfac=1,
# length
train_unet_number=1,
max_length_Y=1, # for Y, X_forImageCondi
max_text_len=1, # for X_forTextCondi
# other info
steps=None,
e=None,
sample_dir=None,
prefix=None,
IF_showfig=True,
CKeys=None,
# TBA to Model B
normalize_X_cond_to_one=False,
# ++
pLM_Model=None, # pLM_Model,
pLM_Model_Name=None, # pLM_Model_Name,
image_channels=None, # image_channels,
pLM_alphabet=None, # esm_alphabet,
):
# -----------
# steps=0
# e=flag
# --
# print (f"Producing {len(X)} samples...")
# ++
if X!=None:
print (f"Producing {len(X)} samples...from text conditioning X...")
lenn_val=len(X)
if X_string!=None:
lenn_val=len(X_string)
print (f"Producing {len(X_string)} samples...from text conditioning X_String (from string)...")
if x_data!=None:
print (f"Producing {len(x_data)} samples...from image conditingig x_data ...")
lenn_val=len(x_data)
# print (x_data)
print ('Device: ', model.device)
pdb_list=[]
fasta_list=[]
for iisample in range (lenn_val):
print(f"Working on {iisample}")
X_cond=None
if X_string==None and X!=None: # for X channel
X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
if X_string!=None: # from raw text, ie., X_string: need tokenizer_X and Xnormfac
# -
# X = tokenizer_X.texts_to_sequences(X_string[iisample])
# X= sequence.pad_sequences(X, maxlen=max_length, padding='post', truncating='post')
# X=np.array(X)
# X_cond=torch.from_numpy(X).float()/Xnormfac
# +
XX = tokenizer_X_forTextCondi.texts_to_sequences(X_string[iisample])
XX = sequence.pad_sequences(XX, maxlen=max_text_len, padding='post', truncating='post')
XX = np.array(XX)
X_cond = torch.from_numpy(XX).float()/Xnormfac_forTextCondi
print ('Tokenized and processed: ', X_cond)
if X_cond!=None:
if normalize_X_cond_to_one: # used when there is constrain on X_cond.sum()
X_cond=X_cond/X_cond.sum()
print ("Text conditoning used: ", X_cond, "...sum: ", X_cond.sum(), "cond scale: ", cond_scales)
else:
print ("Text conditioning used: None")
# for now, assume image_condi and text_condi can be used at the same time
if tokenizer_X_forImageCondi==None:
# ===========================================================
# condi_image/seq needs no tokenization, like numbers: force_path
# only normalization needed
# Based on ModelB:Force_Path
if x_data!=None:
x_data_tokenized=torch.from_numpy(x_data[iisample]/Xnormfac_forImageCondi)
x_data_tokenized=x_data_tokenized.to(torch.float)
# + for debug:
if CKeys['Debug_TrainerPack']==1:
print("x_data_tokenized dim: ", x_data_tokenized.shape)
print("x_data_tokenized dtype: ", x_data_tokenized.dtype)
print("test: ", x_data_tokenized!=None)
else:
x_data_tokenized=None
# + for debug:
if CKeys['Debug_TrainerPack']==1:
print("x_data_tokenized and x_data: None")
# model.sample:full arguments
# self,
# x=None,
# stop_at_unet_number=1,
# cond_scale=7.5,
# # ++
# x_data=None, # image_condi data
# skip_steps=None,
# inpaint_images = None,
# inpaint_masks = None,
# inpaint_resample_times = 5,
# init_images = None,
# x_data_tokenized=None,
# tokenizer_X=None,
# Xnormfac=1.,
# # -+
# device=None,
# max_length=1., # for XandY data, in image/sequence format; NOT for text condition
# max_text_len=1., # for X data, in text format
#
result_embedding=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales ,
x_data=None,
# ++
x_data_tokenized=x_data_tokenized,
#
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,
device=model.device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X_forImageCondi, # tokenizer_X,
Xnormfac=Xnormfac_forImageCondi, # Xnormfac,
# ynormfac=ynormfac,
max_length=max_length_Y, # for ImageCondi, max_length,
max_text_len=max_text_len,
)
else:
# this is for model B in the future
# two channels should be provided: raw cond_img+img_tokenizer or tokenized_cond_img
# need to BE UPDATE and merge with code from
# fun.sample_sequence_omegafold_pLM_ModelB
# one branch is currently missing
result_embedding=model.sample (
x=X_cond,
stop_at_unet_number=train_unet_number ,
cond_scale=cond_scales ,
x_data=x_data[iisample],
# ++
x_data_tokenized=None,
#
skip_steps=skip_steps,
inpaint_images = inpaint_images,
inpaint_masks = inpaint_masks,
inpaint_resample_times = inpaint_resample_times,
init_images = init_images,
device=model.device,
# ++++++++++++++++++++++++++
tokenizer_X=tokenizer_X_forImageCondi, # tokenizer_X,
Xnormfac=Xnormfac_forImageCondi, # Xnormfac,
# ynormfac=ynormfac,
max_length=max_length_Y, # max_length,
max_text_len=max_text_len,
)
# # ------------------------------------------
# result=model.sample (
# X_cond,
# stop_at_unet_number=train_unet_number,
# cond_scale=cond_scales
# )
# # -----------------------------------------------
# result=torch.round(result*ynormfac)
# +++++++++++++++++++++++++++++++++++++++++++++++
# ++ for pLM
# full record
# result_embedding as image.dim: [batch, channels, seq_len]
# result_tokens.dim: [batch, seq_len]
result_tokens,result_logits = convert_into_tokens(
pLM_Model,
result_embedding,
pLM_Model_Name,
)
result=result_tokens.unsqueeze(1) # dim: [batch, 1, seq_len]
# + for debug
print("result.dim: ", result.shape)
fig=plt.figure()
plt.plot (
result[0,0,:].cpu().detach().numpy(),
label= f'Predicted'
)
#plt.plot (GT[samples,0,:]*ynormfac,label= f'GT {0}')
plt.legend()
outname = sample_dir+ f"sampled_from_X_{iisample}_condscale-{str (cond_scales)}_{e}_{steps}.jpg"
#plt.title (f"Sample {samples}, cond scale={str (cond_scales[iisample])}")
if IF_showfig==1:
plt.show ()
else:
plt.savefig(outname, dpi=200)
plt.close()
# # ----------------------------------------
# plt.plot (result[0,0,:].cpu().detach().numpy(),label= f'Predicted')
# plt.legend()
# outname = prefix+ f"sampld_from_X_{flag}_condscale-{str (cond_scales)}_{e}_{steps}.jpg"
# plt.savefig(outname, dpi=200)
# plt.show ()
# # ---------------------------------------------------------
# to_rev=result[:,0,:]
# to_rev=to_rev.long().cpu().detach().numpy()
# print("to_rev.dim: ", to_rev.shape)
# y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
# for iii in range (len(y_data_reversed)):
# y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
to_rev=result[:,0,:]
# the following fun decides ending automatically
y_data_reversed=decode_many_ems_token_rec_for_folding(
to_rev,
result_logits,
pLM_alphabet,
pLM_Model,
)
if CKeys['Debug_TrainerPack']==3:
print("on y_data_reversed[0]: ", y_data_reversed[0])
# + from Model B
### reverse second structure input....
#
if X_cond != None:
# there is condi_text
if X_string!=None:
X_cond=torch.round(X_cond*Xnormfac_forTextCondi)
to_rev=X_cond[:,:]
to_rev=to_rev.long().cpu().detach().numpy()
print ("to_rev.dim: ", to_rev.shape)
# --
# X_data_reversed=tokenizer_X.sequences_to_texts (to_rev)
# ++
X_text_reversed=tokenizer_X_forTextCondi.sequences_to_texts (to_rev)
for iii in range (len(y_text_reversed)):
X_text_reversed[iii]=X_text_reversed[iii].upper().strip().replace(" ", "")
if X_string==None:
# reverse this: X_cond=torch.Tensor (X[iisample]).to(device).unsqueeze (0)
X_text_reversed=X_cond
else:
X_text_reversed=None
if x_data !=None: # there is condi_image
x_data_reversed=x_data #is already in sequence fromat..
else:
x_data_reversed=None
# summary
print (f"For {X_text_reversed} or {X[iisample]} on Text_Condi,\n and {x_data_reversed} on Image_Condi,\n predicted sequence: ", y_data_reversed)
# + for debug
print("================================================")
print("foldproteins: ", foldproteins)
if not foldproteins:
pdb_file=None
else:
# if foldproteins:
if X_cond != None:
pass
tempname='temp'
pdb_file, fasta_file=foldandsavePDB_pdb_fasta (
sequence=y_data_reversed[0],
filename_out=tempname,
num_cycle=num_cycle,
flag=flag,
# +++++++++++++++++++
# prefix=prefix,
prefix=sample_dir,
)
#
out_nam=iisample
#
# out_nam_fasta=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.fasta'
# write_fasta (y_data_reversed[0], out_nam_fasta)
#
out_nam=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.pdb'
out_nam_fasta=f'{sample_dir}DeNovoSampling_{iisample}_epo_{e}_step_{steps}.fasta'
shutil.copy (pdb_file, out_nam) #source, dest
shutil.copy (fasta_file, out_nam_fasta)
# clean the slade to avoid mistakenly using the previous fasta file
os.remove (pdb_file)
os.remove (fasta_file)
#
pdb_file=out_nam
fasta_file=out_nam_fasta
#
pdb_list.append(pdb_file)
fasta_list.append(fasta_file)
#
print (f"Properly named PDB file produced: {pdb_file}")
if IF_showfig==1:
#flag=1000
view=show_pdb(
pdb_file=pdb_file,
flag=flag,
show_sidechains=show_sidechains,
show_mainchains=show_mainchains,
color=color
)
view.show()
if calc_error:
if CKeys['Problem_ID']==7:
# only work for ModelA:SecStr
get_Model_A_error (pdb_file, X[iisample], plotit=True)
else:
print("Error calculation on the predicted results is not applicable...")
return pdb_list, fasta_list
# + TBU
# ++++++++++++++++++++++++++++++++++++++++++++++++
def sample_loop_omegafold_ModelA (
model,
train_loader,
cond_scales=None, # [7.5], #list of cond scales - each sampled...
num_samples=None, # 2, #how many samples produced every time tested.....
timesteps=None, # 100, # not used
flag=None, # 0,
foldproteins=False,
#
cond_image=False, # use_text_embedd=True,
cond_text=True,
skip_steps=0,
#
max_text_len=None,
max_length=None,
# +++++++++++++++++++
train_unet_number=1,
ynormfac=None,
prefix=None,
tokenizer_y=None,
Xnormfac_CondiText=1,
tokenizer_X_CondiText=None,
# ++
CKeys=None,
sample_dir=None,
steps=None,
e=None,
IF_showfig=True, # effective only after foldproteins=True
):
# =====================================================
# sample # = num_samples*(# of mini-batches)
# =====================================================
# steps=0
# e=flag
# for item in train_loader:
for idx, item in enumerate(train_loader):
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
GT=y_train_batch.cpu().detach()
GT= GT.unsqueeze(1)
if num_samples>y_train_batch.shape[0]:
print("Warning: sampling # > len(mini_batch)")
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
X_train_batch_picked = X_train_batch[:num_samples,:]
print ('(TEST) X_batch shape: ', X_train_batch_picked.shape)
# loop over cond_scales:list
for iisample in range (len (cond_scales)):
# ++ for model A
result=model.sample (
x=X_train_batch_picked,
stop_at_unet_number=train_unet_number,
cond_scale=cond_scales[iisample],
#
skip_steps=skip_steps,
device=model.device,
#
max_length=max_length,
max_text_len=max_text_len,
)
# # ++ for model B
# if use_text_embedd:
# result=model.sample (
# # x= X_train_batch,
# x= X_train_batch_picked,
# stop_at_unet_number=train_unet_number ,
# cond_scale=cond_scales[iisample],
# device=device,
# skip_steps=skip_steps
# )
# else:
# result=model.sample (
# x= None,
# # x_data_tokenized= X_train_batch,
# x_data_tokenized= X_train_batch_picked,
# stop_at_unet_number=train_unet_number ,
# cond_scale=cond_scales[iisample],
# device=device,
# skip_steps=skip_steps
# )
result=torch.round(result*ynormfac)
GT=torch.round (GT*ynormfac)
for samples in range (num_samples):
print ("sample ", samples+1, "out of ", num_samples)
fig=plt.figure()
plt.plot (
result[samples,0,:].cpu().detach().numpy(),
label= f'Predicted'
)
plt.plot (
GT[samples,0,:],
label= f'GT {0}'
)
plt.legend()
outname = sample_dir+ f"Batch_{idx}_sample_{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
if IF_showfig==1:
plt.show()
else:
plt.savefig(outname, dpi=200)
plt.close ()
#reverse y sequence
to_rev=result[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
#reverse GT_y sequence
to_rev=GT[:,0,:]
to_rev=to_rev.long().cpu().detach().numpy()
GT_y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
GT_y_data_reversed[iii]=GT_y_data_reversed[iii].upper().strip().replace(" ", "")
### reverse second structure input....
# pay attension to the shape of Xnormfac
# -
# to_rev=torch.round (X_train_batch[:,:]*Xnormfac_CondiText)
# +
to_rev=torch.round (X_train_batch[:,:]*torch.FloatTensor(Xnormfac_CondiText).to(model.device))
to_rev=to_rev.long().cpu().detach().numpy()
# ++ different input
if CKeys['Debug_TrainerPack']==1:
print("tokenizer_X_CondiText: ", tokenizer_X_CondiText)
print("Xnormfac_CondiText: ", Xnormfac_CondiText)
if tokenizer_X_CondiText!=None:
X_data_reversed=tokenizer_X_CondiText.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
else:
X_data_reversed=to_rev.copy()
# + for debug
if CKeys['Debug_TrainerPack']==1:
print("X_data_reversed: ", X_data_reversed)
print (f"For {X_train_batch[samples,:].cpu().detach().numpy()} or {X_data_reversed[samples]}, \npredicted sequence: ", y_data_reversed[samples])
print (f"Ground truth: {GT_y_data_reversed[samples]}")
if foldproteins:
xbc=X_train_batch[samples,:].cpu().detach().numpy()
out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})
tempname='temp'
pdb_file=foldandsavePDB (
sequence=y_data_reversed[samples],
filename_out=tempname,
num_cycle=16, flag=flag,
# +++++++++++++++++++
prefix=prefix
)
# #out_nam=f'{prefix}{out_nam}.pdb'
# out_nam=f'{prefix}{X_data_reversed[samples]}.pdb'
# ------------------------------------------------------
# sometime, this name below can get too long to fit
# out_nam=f'{sample_dir}{X_data_reversed[samples]}.pdb'
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
# add a way to save the sampling name and results
# ref: outname = sample_dir+ f"sample-{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
out_nam=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.pdb'
out_nam_inX=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.txt'
if CKeys['Debug_TrainerPack']==1:
print("pdb_file: ", pdb_file)
print("out_nam: ", out_nam)
print (f'Original PDB: {pdb_file} OUT: {out_nam}')
shutil.copy (pdb_file, out_nam) #source, dest
# +
with open(out_nam_inX, "w") as inX_file:
inX_file.write(f'{X_data_reversed[samples]}\n')
pdb_file=out_nam
print (f"Properly named PDB file produced: {pdb_file}")
print (f"input X for sampling stored: {pdb_file}")
if IF_showfig==1:
view=show_pdb(
pdb_file=pdb_file,
flag=flag,
show_sidechains=show_sidechains,
show_mainchains=show_mainchains,
color=color
)
view.show()
# steps=steps+1
# if steps>num_samples:
# break
# + TBU
# ++++++++++++++++++++++++++++++++++++++++++++++++
def sample_loop_omegafold_pLM_ModelA (
model,
train_loader,
cond_scales=None, # [7.5], #list of cond scales - each sampled...
num_samples=None, # 2, #how many samples produced every time tested.....
timesteps=None, # 100, # not used
flag=None, # 0,
foldproteins=False,
#
cond_image=False, # use_text_embedd=True,
cond_text=True,
skip_steps=0,
#
max_text_len=None,
max_length=None,
# +++++++++++++++++++
train_unet_number=1,
ynormfac=None,
prefix=None,
tokenizer_y=None,
Xnormfac_CondiText=1,
tokenizer_X_CondiText=None,
# ++
CKeys=None,
sample_dir=None,
steps=None,
e=None,
IF_showfig=True, # effective only after foldproteins=True
# ++ for pLM
pLM_Model=None,
pLM_Model_Name=None,
image_channels=None,
pLM_alphabet=None,
# ++ for on-fly check: for SecStr only
calc_error=False, # for check on folded results, not used for every case
):
# =====================================================
# sample # = num_samples*(# of mini-batches)
# =====================================================
# steps=0
# e=flag
# for item in train_loader:
for idx, item in enumerate(train_loader):
X_train_batch= item[0].to(device)
y_train_batch=item[1].to(device)
GT=y_train_batch.cpu().detach()
GT= GT.unsqueeze(1)
if num_samples>y_train_batch.shape[0]:
print("Warning: sampling # > len(mini_batch)")
num_samples = min (num_samples,y_train_batch.shape[0] )
print (f"Producing {num_samples} samples...")
X_train_batch_picked = X_train_batch[:num_samples,:]
print ('(TEST) X_batch shape: ', X_train_batch_picked.shape)
# loop over cond_scales:list
for iisample in range (len (cond_scales)):
# ++ for model A
result_embedding = model.sample (
x=X_train_batch_picked,
stop_at_unet_number=train_unet_number,
cond_scale=cond_scales[iisample],
#
skip_steps=skip_steps,
device=model.device,
#
max_length=max_length,
max_text_len=max_text_len,
#
x_data=None,
x_data_tokenized=None,
#
tokenizer_X=tokenizer_X_CondiText,
Xnormfac=Xnormfac_CondiText,
)
# # ++ for model B
# if use_text_embedd:
# result=model.sample (
# # x= X_train_batch,
# x= X_train_batch_picked,
# stop_at_unet_number=train_unet_number ,
# cond_scale=cond_scales[iisample],
# device=device,
# skip_steps=skip_steps
# )
# else:
# result=model.sample (
# x= None,
# # x_data_tokenized= X_train_batch,
# x_data_tokenized= X_train_batch_picked,
# stop_at_unet_number=train_unet_number ,
# cond_scale=cond_scales[iisample],
# device=device,
# skip_steps=skip_steps
# )
# ++ for pLM:
# full record
# result_embedding as image.dim: [batch, channels, seq_len]
# result_tokens.dim: [batch, seq_len]
result_tokens,result_logits = convert_into_tokens(
pLM_Model,
result_embedding,
pLM_Model_Name,
)
# # --------------------------------------------
# result=torch.round(result*ynormfac)
# GT=torch.round (GT*ynormfac)
# ++++++++++++++++++++++++++++++++++++++++++++
result=result_tokens.unsqueeze(1) # dim: [batch, 1, seq_len]
# # ---------------------------------------------------------
# #reverse y sequence
# to_rev=result[:,0,:]
# to_rev=to_rev.long().cpu().detach().numpy()
# y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
# for iii in range (len(y_data_reversed)):
# y_data_reversed[iii]=y_data_reversed[iii].upper().strip().replace(" ", "")
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
to_rev=result[:,0,:] # token (batch,seq_len)
y_data_reversed=decode_many_ems_token_rec_for_folding(
to_rev,
result_logits,
pLM_alphabet,
pLM_Model,
)
if CKeys['Debug_TrainerPack']==3:
print("on y_data_reversed[0]: ", y_data_reversed[0])
# # -----------------------------------------------------------
# #reverse GT_y sequence
# to_rev=GT[:,0,:]
# to_rev=to_rev.long().cpu().detach().numpy()
# GT_y_data_reversed=tokenizer_y.sequences_to_texts (to_rev)
# for iii in range (len(y_data_reversed)):
# GT_y_data_reversed[iii]=GT_y_data_reversed[iii].upper().strip().replace(" ", "")
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#reverse GT_y sequence
# GT should be SAFE to reverse
to_rev=GT[:,0,:]
GT_y_data_reversed=decode_many_ems_token_rec(
to_rev,
pLM_alphabet,
)
### reverse second structure input....
# pay attension to the shape of Xnormfac
# -
# to_rev=torch.round (X_train_batch[:,:]*Xnormfac_CondiText)
# +
# print("X_train_batch", X_train_batch)
# print("Xnormfac_CondiText: ", Xnormfac_CondiText)
# to_rev=torch.round (X_train_batch[:,:]*torch.FloatTensor(Xnormfac_CondiText).to(model.device))
# print("X_train_batch: ", X_train_batch[:,:])
# print("torch.tensor(Xnormfac_CondiText): ", torch.tensor(Xnormfac_CondiText))
to_rev=X_train_batch[:,:]*torch.tensor(Xnormfac_CondiText).to(model.device)
# print("to_rev ", to_rev)
# # -: convert into int64
# to_rev=to_rev.long().cpu().detach().numpy()
# +: just float
to_rev=to_rev.cpu().detach().numpy()
# print("to_rev 2", to_rev)
# ++ different input
if CKeys['Debug_TrainerPack']==1:
print("tokenizer_X_CondiText: ", tokenizer_X_CondiText)
print("Xnormfac_CondiText: ", Xnormfac_CondiText)
if tokenizer_X_CondiText!=None:
# round the number into tokens
to_rev = np.round(to_rev)
X_data_reversed=tokenizer_X_CondiText.sequences_to_texts (to_rev)
for iii in range (len(y_data_reversed)):
X_data_reversed[iii]=X_data_reversed[iii].upper().strip().replace(" ", "")
else:
X_data_reversed=to_rev.copy()
# + for debug
if CKeys['Debug_TrainerPack']==1:
print("X_data_reversed: ", X_data_reversed)
print("X_data_reversed.dim: ", X_data_reversed.shape)
for samples in range (num_samples):
print ("sample ", samples+1, "out of ", num_samples)
fig=plt.figure()
plt.plot (
result[samples,0,:].cpu().detach().numpy(),
label= f'Predicted'
)
plt.plot (
GT[samples,0,:],
label= f'GT {0}'
)
plt.legend()
outname = sample_dir+ f"Batch_{idx}_sample_{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
if IF_showfig==1:
plt.show()
else:
plt.savefig(outname, dpi=200)
plt.close ()
print (f"For input in dataloader: {X_train_batch[samples,:].cpu().detach().numpy()} or \n recovered input {X_data_reversed[samples]}")
print (f"predicted sequence: {y_data_reversed[samples]}")
print (f"Ground truth: {GT_y_data_reversed[samples]}")
if foldproteins:
# check whether the predicted sequence is valid
if len(y_data_reversed[samples])>0:
# # --
# xbc=X_train_batch[samples,:].cpu().detach().numpy()
# # out_nam=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})
# out_nam_content=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.1f" % xbc})
# ++
xbc=X_data_reversed[samples]
out_nam_content=np.array2string(xbc, formatter={'float_kind':lambda xbc: "%.4f" % xbc})
#
tempname='temp'
pdb_file,fasta_file=foldandsavePDB_pdb_fasta (
sequence=y_data_reversed[samples],
filename_out=tempname,
num_cycle=16, flag=flag,
# +++++++++++++++++++
prefix=prefix
)
# #out_nam=f'{prefix}{out_nam}.pdb'
# out_nam=f'{prefix}{X_data_reversed[samples]}.pdb'
# ------------------------------------------------------
# sometime, this name below can get too long to fit
# out_nam=f'{sample_dir}{X_data_reversed[samples]}.pdb'
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++
# add a way to save the sampling name and results
# ref: outname = sample_dir+ f"sample-{samples}_condscale-{str (cond_scales[iisample])}_{e}_{steps}.jpg"
out_nam=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.pdb'
out_nam_seq=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.fasta'
out_nam_inX=f'{sample_dir}SamplingLoop_B_{idx}_Sample_{samples}_condscale-{str (cond_scales[iisample])}_epo_{e}_step_{steps}.txt'
if CKeys['Debug_TrainerPack']==1:
print("pdb_file: ", pdb_file)
print("out_nam: ", out_nam)
print (f'Original PDB: {pdb_file} OUT: {out_nam}')
shutil.copy (pdb_file, out_nam) #source, dest
shutil.copy (fasta_file, out_nam_seq)
# +
with open(out_nam_inX, "w") as inX_file:
# inX_file.write(f'{X_data_reversed[samples]}\n')
inX_file.write(out_nam_content)
# clean the slade to avoid mistakenly using the previous fasta file
os.remove (pdb_file)
os.remove (fasta_file)
pdb_file=out_nam
print (f"Properly named PDB file produced: {pdb_file}")
print (f"input X for sampling stored: {pdb_file}")
if IF_showfig==1:
view=show_pdb(
pdb_file=pdb_file,
flag=flag,
show_sidechains=show_sidechains,
show_mainchains=show_mainchains,
color=color
)
view.show()
if calc_error:
print('On-fly check...')
if CKeys['Problem_ID']==7:
# only work for ModelA:SecStr
get_Model_A_error (pdb_file, X_data_reversed[samples], plotit=True)
else:
print("Error calculation on the predicted results is not applicable...")
else:
print("The predicted sequence is EMPTY...")