导入库（用于深度学习）
import os import time from datetime import timedelta import json import yaml from tqdm import tqdm import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader
绘图（用于各种，以信号为例）
plt.figure(1) plt.subplot(2, 1, 1) plt.plot(t, sig) # 绘制信号波形图 plt.subplot(2, 1, 2) plt.imshow(np.abs(st_Res), origin='lower', extent=(0, len(t), 0, len(t)//2)) # 绘制频谱图、对复数格式取绝对值即可 plt.savefig("./imgs/stockwell-asdvalve09-" + time.strftime("%Y%m%d%H%M", time.localtime())) plt.show()
深度学习训练模板（torch框架）：
import os import time from datetime import timedelta import json import yaml from tqdm import tqdm import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader class BaseTrainer(object): def __init__(self): self.device = torch.device("cuda") if torch.cuda.is_available() else "cpu" self.model = None self.optimizer = None self.loss = None def __setup_dataloader(self, is_train): if is_train: self.train_dataset = None self.train_loader = DataLoader(self.train_dataset, batch_size=64, shuffle=True, num_workers=0) # 获取测试数据 self.valid_dataset = None self.valid_loader = DataLoader(self.valid_dataset, batch_size=64, shuffle=True, num_workers=0) def __setup_model(self, is_train): self.model = None self.model.to(self.device) # optimizer & scheduler self.loss = .to(self.device) if is_train: if self.configs.train_conf.enable_amp: self.amp_scaler = torch.cuda.amp.GradScaler(init_scale=1024) self.optimizer = torch.optim.Adam(self.model.parameters(), lr=float(self.configs.optimizer_conf.learning_rate)) def __train_epoch(self, epoch_id): pass def train(self): pass def evaluate(self): pass def __test(self, resume_model=None): pass
保存checkpoint、加载checkpoint（torch框架）：
def __load_checkpoint(self, save_model_path, resume_model): last_epoch = -1 best_auc, best_pauc = 0, 0 last_model_dir = os.path.join(save_model_path, f'{self.configs.use_model}', 'last_model') if resume_model is not None or (os.path.exists(os.path.join(last_model_dir, 'model.pth')) and os.path.exists(os.path.join(last_model_dir, 'optimizer.pth'))): # 自动获取最新保存的模型 if resume_model is None: resume_model = last_model_dir assert os.path.exists(os.path.join(resume_model, 'model.pth')), "模型参数文件不存在！" assert os.path.exists(os.path.join(resume_model, 'optimizer.pth')), "优化方法参数文件不存在！" state_dict = torch.load(os.path.join(resume_model, 'model.pth')) if isinstance(self.model, torch.nn.parallel.DistributedDataParallel): self.model.module.load_state_dict(state_dict) else: self.model.load_state_dict(state_dict) self.optimizer.load_state_dict(torch.load(os.path.join(resume_model, 'optimizer.pth'))) # 自动混合精度参数 if self.amp_scaler is not None and os.path.exists(os.path.join(resume_model, 'scaler.pth')): self.amp_scaler.load_state_dict(torch.load(os.path.join(resume_model, 'scaler.pth'))) with open(os.path.join(resume_model, 'model.state'), 'r', encoding='utf-8') as f: json_data = json.load(f) last_epoch = json_data['last_epoch'] - 1 best_auc = json_data['best_auc'] best_pauc = json_data['best_pauc'] self.logger.info('成功恢复模型参数和优化方法参数：{}'.format(resume_model)) self.optimizer.step() [self.scheduler.step() for _ in range(last_epoch * len(self.train_loader))] return last_epoch, best_auc, best_pauc def __save_checkpoint(self, save_model_path, epoch_id, best_auc=0., best_pauc=0., best_model=False): if isinstance(self.model, torch.nn.parallel.DistributedDataParallel): state_dict = self.model.module.state_dict() else: state_dict = self.model.state_dict() if best_model: model_path = os.path.join(save_model_path, f'{self.configs.use_model}', 'best_model') else: model_path = os.path.join(save_model_path, f'{self.configs.use_model}', 'epoch_{}'.format(epoch_id)) os.makedirs(model_path, exist_ok=True) torch.save(self.optimizer.state_dict(), os.path.join(model_path, 'optimizer.pth')) torch.save(state_dict, os.path.join(model_path, 'model.pth')) # 自动混合精度参数 if self.amp_scaler is not None: torch.save(self.amp_scaler.state_dict(), os.path.join(model_path, 'scaler.pth')) with open(os.path.join(model_path, 'model.state'), 'w', encoding='utf-8') as f: data = {"last_epoch": epoch_id, "best_auc": best_auc, "best_pauc": best_pauc, "version": __version__} f.write(json.dumps(data)) if not best_model: last_model_path = os.path.join(save_model_path, f'{self.configs.use_model}', 'last_model') shutil.rmtree(last_model_path, ignore_errors=True) shutil.copytree(model_path, last_model_path) # 删除旧的模型 old_model_path = os.path.join(save_model_path, f'{self.configs.use_model}', 'epoch_{}'.format(epoch_id - 3)) if os.path.exists(old_model_path): shutil.rmtree(old_model_path) self.logger.info('已保存模型：{}'.format(model_path))
加载预训练模型（torch）：
def __load_pretrained(self, pretrained_model): # 加载预训练模型 if pretrained_model is not None: if os.path.isdir(pretrained_model): pretrained_model = os.path.join(pretrained_model, 'model.pth') assert os.path.exists(pretrained_model), f"{pretrained_model} 模型不存在！" if isinstance(self.model, torch.nn.parallel.DistributedDataParallel): model_dict = self.model.module.state_dict() else: model_dict = self.model.state_dict() model_state_dict = torch.load(pretrained_model) # 过滤不存在的参数 for name, weight in model_dict.items(): if name in model_state_dict.keys(): if list(weight.shape) != list(model_state_dict[name].shape): self.logger.warning('{} not used, shape {} unmatched with {} in model.'. format(name, list(model_state_dict[name].shape), list(weight.shape))) model_state_dict.pop(name, None) else: self.logger.warning('Lack weight: {}'.format(name)) if isinstance(self.model, torch.nn.parallel.DistributedDataParallel): self.model.module.load_state_dict(model_state_dict, strict=False) else: self.model.load_state_dict(model_state_dict, strict=False) self.logger.info('成功加载预训练模型：{}'.format(pretrained_model))