| import traceback
|
| from collections import OrderedDict
|
| from time import time as ttime
|
| import shutil,os
|
| import torch
|
| from tools.i18n.i18n import I18nAuto
|
|
|
| i18n = I18nAuto()
|
|
|
| def my_save(fea,path):
|
| dir=os.path.dirname(path)
|
| name=os.path.basename(path)
|
| tmp_path="%s.pth"%(ttime())
|
| torch.save(fea,tmp_path)
|
| shutil.move(tmp_path,"%s/%s"%(dir,name))
|
|
|
| def savee(ckpt, name, epoch, steps, hps):
|
| try:
|
| opt = OrderedDict()
|
| opt["weight"] = {}
|
| for key in ckpt.keys():
|
| if "enc_q" in key:
|
| continue
|
| opt["weight"][key] = ckpt[key].half()
|
| opt["config"] = hps
|
| opt["info"] = "%sepoch_%siteration" % (epoch, steps)
|
|
|
| my_save(opt, "%s/%s.pth" % (hps.save_weight_dir, name))
|
| return "Success."
|
| except:
|
| return traceback.format_exc()
|
|
|
