Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,457 Bytes
eadd7b4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import os
import re
import torch
from diffusion.utils.logger import get_root_logger
def save_checkpoint(work_dir,
epoch,
model,
model_ema=None,
optimizer=None,
lr_scheduler=None,
keep_last=False,
step=None,
):
os.makedirs(work_dir, exist_ok=True)
state_dict = dict(state_dict=model.state_dict())
if model_ema is not None:
state_dict['state_dict_ema'] = model_ema.state_dict()
if optimizer is not None:
state_dict['optimizer'] = optimizer.state_dict()
if lr_scheduler is not None:
state_dict['scheduler'] = lr_scheduler.state_dict()
if epoch is not None:
state_dict['epoch'] = epoch
file_path = os.path.join(work_dir, f"epoch_{epoch}.pth")
if step is not None:
file_path = file_path.split('.pth')[0] + f"_step_{step}.pth"
logger = get_root_logger()
torch.save(state_dict, file_path)
logger.info(f'Saved checkpoint of epoch {epoch} to {file_path.format(epoch)}.')
if keep_last:
for i in range(epoch):
previous_ckgt = file_path.format(i)
if os.path.exists(previous_ckgt):
os.remove(previous_ckgt)
def load_checkpoint(checkpoint,
model,
model_ema=None,
optimizer=None,
lr_scheduler=None,
load_ema=False,
resume_optimizer=True,
resume_lr_scheduler=True,
max_length=120,
):
assert isinstance(checkpoint, str)
ckpt_file = checkpoint
checkpoint = torch.load(ckpt_file, map_location="cpu")
state_dict_keys = ['pos_embed', 'base_model.pos_embed', 'model.pos_embed']
for key in state_dict_keys:
if key in checkpoint['state_dict']:
del checkpoint['state_dict'][key]
if 'state_dict_ema' in checkpoint and key in checkpoint['state_dict_ema']:
del checkpoint['state_dict_ema'][key]
break
if load_ema:
state_dict = checkpoint['state_dict_ema']
else:
state_dict = checkpoint.get('state_dict', checkpoint) # to be compatible with the official checkpoint
null_embed = torch.load(f'output/pretrained_models/null_embed_diffusers_{max_length}token.pth', map_location='cpu')
state_dict['y_embedder.y_embedding'] = null_embed['uncond_prompt_embeds'][0]
missing, unexpect = model.load_state_dict(state_dict, strict=False)
if model_ema is not None:
model_ema.load_state_dict(checkpoint['state_dict_ema'], strict=False)
if optimizer is not None and resume_optimizer:
optimizer.load_state_dict(checkpoint['optimizer'])
if lr_scheduler is not None and resume_lr_scheduler:
lr_scheduler.load_state_dict(checkpoint['scheduler'])
logger = get_root_logger()
if optimizer is not None:
epoch = checkpoint.get('epoch', re.match(r'.*epoch_(\d*).*.pth', ckpt_file).group()[0])
logger.info(f'Resume checkpoint of epoch {epoch} from {ckpt_file}. Load ema: {load_ema}, '
f'resume optimizer： {resume_optimizer}, resume lr scheduler: {resume_lr_scheduler}.')
return epoch, missing, unexpect
logger.info(f'Load checkpoint from {ckpt_file}. Load ema: {load_ema}.')
return missing, unexpect
|