#################################### # # # import libraries # # # #################################### import torch import numpy as np import random import os import json import logging import pandas as pd import torch.distributed as dist from torch.utils.data import Dataset, DataLoader import h5py from transformers import AutoTokenizer from torch.utils.data.distributed import DistributedSampler from functools import partial import torch.nn.functional as F #################################### # # # helper functions # # # #################################### def ensure_dir(path: str) -> None: os.makedirs(path, exist_ok=True) def create_logger(log_file: str, rank: int, name: str = "ProRiboGen") -> logging.Logger: logger = logging.getLogger(name) logger.setLevel(logging.INFO) logger.propagate=False if getattr(logger, "_proribogen_inited", False): return logger fmt = logging.Formatter("[%(asctime)s] [%(levelname)s] %(message)s", "%Y-%m-%d %H:%M:%S") if rank == 0: ensure_dir(os.path.dirname(log_file) or ".") fh = logging.FileHandler(log_file, encoding="utf-8") fh.setFormatter(fmt) fh.setLevel(logging.INFO) logger.addHandler(fh) sh = logging.StreamHandler() sh.setFormatter(fmt) sh.setLevel(logging.INFO) logger.addHandler(sh) else: sh = logging.StreamHandler(open(os.devnull, "w")) sh.setFormatter(fmt) sh.setLevel(logging.CRITICAL) logger.addHandler(sh) logger._proribogen_inited = True return logger def cleanup_ddp(): torch.cuda.empty_cache() if dist.is_available() and dist.is_initialized(): dist.destroy_process_group() def setup_ddp(rank: int, world_size: int, port: int): torch.cuda.set_device(rank) if world_size > 1: os.environ["MASTER_PORT"] = str(port) os.environ["MASTER_ADDR"] = "127.0.0.1" dist.init_process_group( backend="nccl", rank=rank, world_size=world_size ) def is_main_process(rank): return rank == 0 def load_config(config_path: str = "config.json") -> dict: with open(config_path, "r", encoding="utf-8") as f: config = json.load(f) return config def get_h5_embedding_dim(h5_path: str) -> int: """Return per-token protein embedding dimension from a merged embeddings H5.""" with h5py.File(h5_path, "r") as f: if "embeddings" not in f: raise ValueError(f"H5 missing 'embeddings' dataset: {h5_path}") return int(f["embeddings"].shape[1]) def validate_protein_dim_against_h5(config: dict, base_dir: str = ".") -> int: """ Ensure config model.protein_dim matches the H5 embedding width. If protein_dim is omitted, set it from the H5 file. """ h5_rel = config["data"]["protein_h5"] h5_path = h5_rel if os.path.isabs(h5_rel) else os.path.join(base_dir, h5_rel) if not os.path.isfile(h5_path): raise FileNotFoundError(f"Protein H5 not found: {h5_path}") h5_dim = get_h5_embedding_dim(h5_path) model_cfg = config.setdefault("model", {}) cfg_dim = model_cfg.get("protein_dim") if cfg_dim is None: model_cfg["protein_dim"] = h5_dim return h5_dim cfg_dim = int(cfg_dim) if cfg_dim != h5_dim: raise ValueError( f"config model.protein_dim={cfg_dim} does not match H5 embedding dim={h5_dim} " f"({h5_path})" ) return h5_dim def a_useful_log(logger, rank, message): if is_main_process(rank): logger.info(message) #################################### # # # Data # # # #################################### class RnaRbpDataset(Dataset): """ 读取 RNA–RBP（蛋白序列嵌入）配对数据，用于建模 RBP 与靶 RNA 的相互作用。 CSV 主列为 ``r_id``, ``rna``, ``p_id``；若数据仍为旧表头 ``d_id`` / ``dna``，会自动映射到 上述列。 H5: - embeddings : [N_total_tokens, protein_dim] - p_ids : [N_proteins] - starts : [N_proteins] - lengths : [N_proteins] 返回单样本 dict: { "r_id": str, "rna": str, "p_id": str, "rna_len": int, "protein_emb": torch.FloatTensor [Lp, Dp], "protein_len": int, } """ REQUIRED_COLUMNS = ("r_id", "rna", "p_id") _LEGACY_ALIASES = { "r_id": ("d_id",), "rna": ("dna",), } # CSV p_id -> H5 p_id when naming differs between annotation and embedding index _P_ID_ALIASES = { "human-DDX3": "human-DDX3X", "human-HNRNPH1": "human-HNRNPH", } def __init__(self, csv_path: str, h5_path: str = "data/esm2-protein_embeddings.h5") -> None: super().__init__() self.csv_path = csv_path self.h5_path = h5_path self.data = pd.read_csv(csv_path) for canonical, aliases in self._LEGACY_ALIASES.items(): if canonical not in self.data.columns: for alt in aliases: if alt in self.data.columns: self.data[canonical] = self.data[alt] break missing_cols = [col for col in self.REQUIRED_COLUMNS if col not in self.data.columns] if missing_cols: raise ValueError( f"CSV 文件缺少必要列: {missing_cols}. 当前列为: {list(self.data.columns)}" ) self.data["rna"] = self.data["rna"].astype(str).str.upper() self.data["r_id"] = self.data["r_id"].astype(str) self.data["p_id"] = self.data["p_id"].astype(str).replace(self._P_ID_ALIASES) self.data["rna_len"] = self.data["rna"].str.len() with h5py.File(self.h5_path, "r") as f: p_ids = [ x.decode("utf-8") if isinstance(x, bytes) else str(x) for x in f["p_ids"][:] ] starts = f["starts"][:] lengths = f["lengths"][:] self.protein_index = { p_id: (int(start), int(length)) for p_id, start, length in zip(p_ids, starts, lengths) } self.all_p_ids = list(self.protein_index.keys()) valid_mask = self.data["p_id"].isin(self.protein_index) n_dropped = int((~valid_mask).sum()) if n_dropped: missing_p_ids = sorted(self.data.loc[~valid_mask, "p_id"].unique().tolist()) self.data = self.data.loc[valid_mask].reset_index(drop=True) logging.getLogger("ProRiboGen").warning( "Dropped %d rows from %s: p_id not in H5 (%s)", n_dropped, csv_path, ", ".join(missing_p_ids[:20]) + (" ..." if len(missing_p_ids) > 20 else ""), ) self._h5_file = None self._embeddings_ds = None def _ensure_h5_open(self): if self._h5_file is None: self._h5_file = h5py.File(self.h5_path, "r") self._embeddings_ds = self._h5_file["embeddings"] def __len__(self) -> int: return len(self.data) def get_protein_embedding_by_pid(self, p_id: str) -> torch.Tensor: if p_id not in self.protein_index: raise KeyError(f"Protein id not found in H5 index: {p_id}") start, length = self.protein_index[p_id] self._ensure_h5_open() protein_emb = self._embeddings_ds[start : start + length] return torch.from_numpy(protein_emb).float() def __getitem__(self, idx: int) -> dict: row = self.data.iloc[idx] r_id = row["r_id"] rna = row["rna"] p_id = row["p_id"] rna_len = int(row["rna_len"]) start, length = self.protein_index[p_id] self._ensure_h5_open() protein_emb = self._embeddings_ds[start : start + length] protein_emb = torch.from_numpy(protein_emb).float() sample = { "r_id": r_id, "rna": rna, "p_id": p_id, "rna_len": rna_len, "protein_emb": protein_emb, "protein_len": int(length), } return sample def __del__(self): if self._h5_file is not None: try: self._h5_file.close() except Exception: pass def rna_rbp_collate_fn( batch, tokenizer, *, append_eos: bool = False, max_rna_nt: int | None = None, ): r_ids = [sample["r_id"] for sample in batch] p_ids = [sample["p_id"] for sample in batch] rna_strs = [sample["rna"] for sample in batch] if append_eos and getattr(tokenizer, "eos_token", None): eos = tokenizer.eos_token rna_strs = [s + eos for s in rna_strs] rna_lens = torch.tensor([sample["rna_len"] for sample in batch], dtype=torch.long) protein_lens = torch.tensor([sample["protein_len"] for sample in batch], dtype=torch.long) tok_kw: dict = {"padding": True, "return_tensors": "pt"} if max_rna_nt is not None and max_rna_nt > 0: tok_kw["max_length"] = max_rna_nt + 2 tok_kw["truncation"] = True rna_enc = tokenizer(rna_strs, **tok_kw) rna_input_ids_clean = rna_enc["input_ids"] rna_attention_mask = rna_enc["attention_mask"] batch_size = len(batch) max_protein_len = max(sample["protein_emb"].shape[0] for sample in batch) protein_dim = batch[0]["protein_emb"].shape[1] protein_cond = torch.zeros( batch_size, max_protein_len, protein_dim, dtype=batch[0]["protein_emb"].dtype, ) protein_attention_mask = torch.zeros( batch_size, max_protein_len, dtype=torch.long, ) for i, sample in enumerate(batch): protein_emb = sample["protein_emb"] Lp = protein_emb.shape[0] protein_cond[i, :Lp] = protein_emb protein_attention_mask[i, :Lp] = 1 return { "r_id": r_ids, "p_id": p_ids, "rna_input_ids_clean": rna_input_ids_clean, "rna_attention_mask": rna_attention_mask, "rna_len_bp": rna_lens, "protein_cond": protein_cond, "protein_attention_mask": protein_attention_mask, "protein_len_tokens": protein_lens, } def build_dataloaders(config, world_size, rank): num_workers = config["train"]["num_workers"] use_worker_processes = num_workers > 0 # 1. tokenizer tokenizer = AutoTokenizer.from_pretrained( config["data"]["tokenizer_path"], trust_remote_code=True, ) # 2. dataset train_dataset = RnaRbpDataset( csv_path=config["data"]["train_csv"], h5_path=config["data"]["protein_h5"], ) test_dataset = RnaRbpDataset( csv_path=config["data"]["test_csv"], h5_path=config["data"]["protein_h5"], ) append_eos = bool(config.get("data", {}).get("append_eos_token", False)) max_rna_nt = config.get("data", {}).get("max_generated_rna_bp") if max_rna_nt is not None: max_rna_nt = int(max_rna_nt) collate_fn = partial( rna_rbp_collate_fn, tokenizer=tokenizer, append_eos=append_eos, max_rna_nt=max_rna_nt, ) # 4. sampler if world_size > 1: train_sampler = DistributedSampler( train_dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False, ) test_sampler = DistributedSampler( test_dataset, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False, ) train_shuffle = False else: train_sampler = None test_sampler = None train_shuffle = True # 5. dataloader train_loader = DataLoader( train_dataset, batch_size=config["train"]["batch_size_per_gpu"], shuffle=train_shuffle, sampler=train_sampler, num_workers=num_workers, pin_memory=True, drop_last=False, collate_fn=collate_fn, persistent_workers=use_worker_processes, multiprocessing_context="spawn" if use_worker_processes else None, ) test_loader = DataLoader( test_dataset, batch_size=config["train"]["eval_batch_size_per_gpu"], shuffle=False, sampler=test_sampler, num_workers=num_workers, pin_memory=True, drop_last=False, collate_fn=collate_fn, persistent_workers=use_worker_processes, multiprocessing_context="spawn" if use_worker_processes else None, ) return train_loader, test_loader, train_sampler, test_sampler, tokenizer def sample_random_negative_sources( pos_p_ids, all_p_ids, use_batch_negatives_first: bool = True, rng=None, deterministic: bool = False, ): if rng is None: rng = random if len(all_p_ids) <= 1: raise ValueError("Need at least two proteins to sample negatives.") negative_sources = [] batch_p_ids = list(pos_p_ids) for i, pos_p_id in enumerate(batch_p_ids): batch_candidates = [ j for j, candidate_p_id in enumerate(batch_p_ids) if candidate_p_id != pos_p_id ] if use_batch_negatives_first and batch_candidates: chosen_idx = batch_candidates[0] if deterministic else rng.choice(batch_candidates) negative_sources.append( { "source": "batch", "batch_index": chosen_idx, "neg_p_id": batch_p_ids[chosen_idx], } ) continue if deterministic: neg_p_id = next(candidate for candidate in all_p_ids if candidate != pos_p_id) else: neg_p_id = rng.choice(all_p_ids) while neg_p_id == pos_p_id: neg_p_id = rng.choice(all_p_ids) negative_sources.append( { "source": "pool", "neg_p_id": neg_p_id, } ) return negative_sources def build_random_negative_protein_batch( batch: dict, dataset: RnaRbpDataset, use_batch_negatives_first: bool = True, rng=None, deterministic: bool = False, ) -> dict: batch_p_ids = list(batch["p_id"]) negative_sources = sample_random_negative_sources( pos_p_ids=batch_p_ids, all_p_ids=dataset.all_p_ids, use_batch_negatives_first=use_batch_negatives_first, rng=rng, deterministic=deterministic, ) negative_embeddings = [] negative_p_ids = [] negative_lengths = [] batch_protein_cond = batch["protein_cond"] batch_protein_attention_mask = batch["protein_attention_mask"] for source in negative_sources: neg_p_id = source["neg_p_id"] if source["source"] == "batch": batch_index = source["batch_index"] protein_len = int(batch_protein_attention_mask[batch_index].sum().item()) negative_emb = batch_protein_cond[batch_index, :protein_len].clone() else: negative_emb = dataset.get_protein_embedding_by_pid(neg_p_id) protein_len = int(negative_emb.shape[0]) negative_embeddings.append(negative_emb) negative_p_ids.append(neg_p_id) negative_lengths.append(protein_len) batch_size = len(negative_embeddings) max_negative_len = max(negative_lengths) protein_dim = negative_embeddings[0].shape[1] negative_protein_cond = torch.zeros( batch_size, max_negative_len, protein_dim, dtype=negative_embeddings[0].dtype, ) negative_protein_attention_mask = torch.zeros( batch_size, max_negative_len, dtype=torch.long, ) for i, negative_emb in enumerate(negative_embeddings): protein_len = negative_emb.shape[0] negative_protein_cond[i, :protein_len] = negative_emb negative_protein_attention_mask[i, :protein_len] = 1 return { "neg_p_id": negative_p_ids, "neg_protein_cond": negative_protein_cond, "neg_protein_attention_mask": negative_protein_attention_mask, "neg_protein_len_tokens": torch.tensor(negative_lengths, dtype=torch.long), } #################################### # # # Other # # # #################################### base_config = { "add_bias_fnn": False, "attention_probs_dropout_prob": 0.0, "emb_layer_norm_before": False, "esmfold_config": None, "hidden_dropout_prob": 0.0, "hidden_size": 512, "initializer_range": 0.02, "intermediate_size": 2048, "is_folding_model": False, "layer_norm_eps": 1e-12, "mask_token_id": 2, "max_position_embeddings": 2050, "model_type": "esm", "num_attention_heads": 16, "num_hidden_layers": 12, "pad_token_id": 1, "position_embedding_type": "rotary", "tie_word_embeddings": False, "token_dropout": False, "torch_dtype": "float32", "transformers_version": "4.54.1", "use_cache": False, "vocab_list": None, "vocab_size": 4107 } def sample_t( batch_size: int, device: torch.device, eps: float = 1e-3, dtype: torch.dtype = torch.float32,) -> torch.Tensor: """ Sample diffusion corruption ratio t for each sample. Args: batch_size: number of samples in the batch. device: target device. eps: lower bound to avoid t=0. dtype: dtype of the returned tensor. Returns: t: shape [B], each element sampled from Uniform(eps, 1.0). """ if not (0.0 <= eps < 1.0): raise ValueError(f"`eps` must satisfy 0 <= eps < 1, got {eps}.") t = eps + (1.0 - eps) * torch.rand(batch_size, device=device, dtype=dtype) return t def build_xt_and_labels( x0: torch.LongTensor, rna_attention_mask: torch.Tensor, mask_token_id: int, t: torch.Tensor, special_token_ids=None, force_at_least_one_mask: bool = True, ): """ Build corrupted input x_t and MLM-style labels for ProRiboGen training. Args: x0: Clean RNA token ids, shape [B, L]. rna_attention_mask: RNA attention mask, shape [B, L]. Valid tokens = 1, padding = 0. mask_token_id: Token id for . t: Per-sample corruption ratio, shape [B]. special_token_ids: Iterable of token ids that should never be masked, e.g. [pad_id,cls_id]. force_at_least_one_mask: If True, ensure each sample has at least one masked valid token. Returns: xt: Corrupted input ids, shape [B, L]. labels: Target ids for CE loss, shape [B, L], with non-masked positions setto -100. masked_positions: Bool tensor, shape [B, L], True where the token was masked. """ if x0.dim() != 2: raise ValueError(f"`x0` must have shape [B, L], got{tuple(x0.shape)}.") if rna_attention_mask.shape != x0.shape: raise ValueError( f"`rna_attention_mask` must match `x0` shape, got " f"{tuple(rna_attention_mask.shape)} vs {tuple(x0.shape)}." ) if t.dim() != 1 or t.shape[0] != x0.shape[0]: raise ValueError( f"`t` must have shape [B], got {tuple(t.shape)} for batch size{x0.shape[0]}." ) B, L = x0.shape device = x0.device valid_positions = rna_attention_mask.bool() if special_token_ids is not None: for token_id in special_token_ids: if token_id is not None: valid_positions &= (x0 != token_id) mask_probs = t.unsqueeze(1).expand(B, L) random_vals = torch.rand(B, L, device=device) masked_positions = (random_vals < mask_probs) & valid_positions if force_at_least_one_mask: no_mask_rows = masked_positions.sum(dim=1) == 0 if no_mask_rows.any(): no_mask_indices = no_mask_rows.nonzero(as_tuple=True)[0] for b in no_mask_indices.tolist(): candidate_positions = valid_positions[b].nonzero(as_tuple=True)[0] if candidate_positions.numel() == 0: raise ValueError( f"Sample {b} has no valid RNA token available for masking." ) chosen = candidate_positions[ torch.randint( low=0, high=candidate_positions.numel(), size=(1,), device=device, ) ] masked_positions[b, chosen] = True xt = x0.clone() xt[masked_positions] = mask_token_id labels = x0.clone() labels[~masked_positions] = -100 return xt, labels, masked_positions #################################### # # # loss # # # #################################### def compute_diffusion_loss( logits: torch.Tensor, labels: torch.LongTensor, t: torch.Tensor, eps: float = 1e-3,) -> torch.Tensor: if logits.dim() != 3: raise ValueError(f"`logits` must have shape [B, L, V], got{tuple(logits.shape)}.") if labels.dim() != 2: raise ValueError(f"`labels` must have shape [B, L], got{tuple(labels.shape)}.") if logits.shape[:2] != labels.shape: raise ValueError( f"`logits` and `labels` must match on [B, L], got " f"{tuple(logits.shape[:2])} vs {tuple(labels.shape)}.") if t.dim() != 1 or t.shape[0] != logits.shape[0]: raise ValueError( f"`t` must have shape [B], got {tuple(t.shape)} for batch size{logits.shape[0]}.") B, L, V = logits.shape token_loss = F.cross_entropy( logits.float().reshape(B * L, V), labels.view(B * L), reduction="none", ignore_index=-100, ).view(B, L) masked_positions = (labels != -100) sample_loss = (token_loss * masked_positions).sum(dim=1) weights = 1.0 / t.clamp_min(eps) loss = (sample_loss * weights).mean() return loss def compute_per_sample_masked_token_ce( logits: torch.Tensor, labels: torch.LongTensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if logits.dim() != 3: raise ValueError(f"`logits` must have shape [B, L, V], got{tuple(logits.shape)}.") if labels.dim() != 2: raise ValueError(f"`labels` must have shape [B, L], got{tuple(labels.shape)}.") if logits.shape[:2] != labels.shape: raise ValueError( f"`logits` and `labels` must match on [B, L], got " f"{tuple(logits.shape[:2])} vs {tuple(labels.shape)}." ) B, L, V = logits.shape token_loss = F.cross_entropy( logits.float().reshape(B * L, V), labels.view(B * L), reduction="none", ignore_index=-100, ).view(B, L) masked_positions = (labels != -100) sample_loss_sum = (token_loss * masked_positions).sum(dim=1) sample_masked_count = masked_positions.sum(dim=1) sample_masked_ce = sample_loss_sum / sample_masked_count.clamp_min(1) return sample_masked_ce, sample_loss_sum, sample_masked_count def compute_masked_token_ce( logits: torch.Tensor, labels: torch.LongTensor, ) -> tuple[torch.Tensor, torch.Tensor]: _, sample_loss_sum, sample_masked_count = compute_per_sample_masked_token_ce( logits=logits, labels=labels, ) masked_token_loss_sum = sample_loss_sum.sum() masked_token_count = sample_masked_count.sum() masked_token_ce = masked_token_loss_sum / masked_token_count.clamp_min(1) return masked_token_ce, masked_token_count def compute_pn_losses( logits_pos: torch.Tensor, logits_neg: torch.Tensor, labels: torch.LongTensor, t: torch.Tensor, alpha: float, margin: float, diffusion_eps: float = 1e-3, ) -> dict: diff_loss_pos = compute_diffusion_loss( logits=logits_pos, labels=labels, t=t, eps=diffusion_eps, ) pos_sample_ce, pos_sample_loss_sum, pos_sample_masked_count = compute_per_sample_masked_token_ce( logits=logits_pos, labels=labels, ) neg_sample_ce, neg_sample_loss_sum, neg_sample_masked_count = compute_per_sample_masked_token_ce( logits=logits_neg, labels=labels, ) rank_per_sample = torch.relu(margin + pos_sample_ce - neg_sample_ce) rank_loss = rank_per_sample.mean() total_loss = diff_loss_pos + float(alpha) * rank_loss sample_count = torch.tensor( pos_sample_ce.shape[0], device=logits_pos.device, dtype=torch.float64, ) delta_ce_sum = (neg_sample_ce - pos_sample_ce).to(torch.float64).sum() margin_satisfied_sum = ( (neg_sample_ce >= (pos_sample_ce + float(margin))).to(torch.float64).sum() ) return { "diff_loss_pos": diff_loss_pos, "rank_loss": rank_loss, "total_loss": total_loss, "pos_masked_token_loss_sum": pos_sample_loss_sum.to(torch.float64).sum(), "pos_masked_token_count": pos_sample_masked_count.to(torch.float64).sum(), "neg_masked_token_loss_sum": neg_sample_loss_sum.to(torch.float64).sum(), "neg_masked_token_count": neg_sample_masked_count.to(torch.float64).sum(), "delta_ce_sum": delta_ce_sum, "margin_satisfied_sum": margin_satisfied_sum, "rank_loss_sum": rank_per_sample.to(torch.float64).sum(), "sample_count": sample_count, } #################################### # # # save/resume # # # #################################### def save_model( save_path: str, model, optimizer=None, scheduler=None, epoch: int = 0, global_step: int = 0, config: dict = None, extra_state: dict = None,) -> None: """ Save training checkpoint. Args: save_path: checkpoint file path, e.g. "checkpoints/latest.pt" model: model or DDP-wrapped model optimizer: optional optimizer scheduler: optional scheduler epoch: current finished epoch index global_step: current optimizer step count config: optional config dict extra_state: optional extra metadata dict """ ensure_dir(os.path.dirname(save_path) or ".") model_to_save = model.module if hasattr(model, "module") else model checkpoint = { "model_state_dict": model_to_save.state_dict(), "optimizer_state_dict": optimizer.state_dict() if optimizer is not None else None, "scheduler_state_dict": scheduler.state_dict() if scheduler is not None else None, "epoch": epoch, "global_step": global_step, "config": config, "extra_state": extra_state if extra_state is not None else {},} torch.save(checkpoint, save_path) def _remap_legacy_checkpoint_state_dict(state_dict: dict) -> dict: """ 旧 checkpoint 中 ``protein_conditioning_attention.dna_norm.*`` 与当前模块名 ``rna_norm`` 对齐。 """ out = {} for key, value in state_dict.items(): new_key = key.replace( "protein_conditioning_attention.dna_norm.", "protein_conditioning_attention.rna_norm.", ) out[new_key] = value return out def resume_model( load_path: str, model, optimizer=None, scheduler=None, device: torch.device = None, strict: bool = True,) -> dict: """ Resume training checkpoint. Args: load_path: checkpoint file path model: model or DDP-wrapped model optimizer: optional optimizer scheduler: optional scheduler device: target device for optimizer state tensors strict: passed to model.load_state_dict() Returns: A dict with: - epoch - global_step - config - extra_state """ if not os.path.exists(load_path): raise FileNotFoundError(f"Checkpoint not found: {load_path}") checkpoint = torch.load(load_path, map_location="cpu") model_to_load = model.module if hasattr(model, "module") else model state_dict = _remap_legacy_checkpoint_state_dict(checkpoint["model_state_dict"]) model_to_load.load_state_dict(state_dict, strict=strict) if optimizer is not None and checkpoint.get("optimizer_state_dict") is not None: optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) if device is not None: for state in optimizer.state.values(): for k, v in state.items(): if torch.is_tensor(v): state[k] = v.to(device) if scheduler is not None and checkpoint.get("scheduler_state_dict") is not None: scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) resume_state = { "epoch": checkpoint.get("epoch", -1), "global_step": checkpoint.get("global_step", 0), "config": checkpoint.get("config", None), "extra_state": checkpoint.get("extra_state", {}), } return resume_state