Python 代码实现了一个基于 PyTorch 的音频故障检测模型的训练和评估流程
想和大神交朋友或想软件开发兼职接项目,请通过手机端搜小#程#序: "黄页小艺” 。
import argparse
import os
import sys
import random
import shutil
import time
import warnings
from copy import deepcopy
import librosa
import numpy as np
import torch
import torchaudio
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
import torch.multiprocessing as mp
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
from torch.utils.data import DataLoader
from mqbench.convert_deploy import convert_deploy
from mqbench.prepare_by_platform import prepare_by_platform, BackendType
from mqbench.utils.state import enable_calibration, enable_quantization, disable_all
from tensorboardX import SummaryWriter
# import my model
from models import cnn
from MachineDataset import MachineDataset
parser = argparse.ArgumentParser(description='PyTorch MIMII Training')
parser.add_argument('--dataset', default='E:\\semiconductor_institute\\MachineFaultDetection\\DataSet\\fan', metavar='DIR',
help='path to dataset', required=False)
'''
parser.add_argument('--train_data', default='E:\\semiconductor_institute\\MachineFaultDetection\\DataSet\\fan_MQB\\fan', metavar='DIR',
help='path to dataset', required=False)
parser.add_argument('--val_data', default='E:\\semiconductor_institute\\MachineFaultDetection\\DataSet\\fan_MQB\\fan', metavar='DIR',
help='path to dataset', required=False)
'''
parser.add_argument('--audio_height', default=128, type=int, metavar='N',
help='image height of training')
parser.add_argument('--audio_width', default=1001, type=int, metavar='N',
help='image width of training')
parser.add_argument('-a', '--arch', default='cnn',metavar='ARCH',
help='model architecture: ')
parser.add_argument('-j', '--workers', default=1, type=int, metavar='N',
help='number of data loading workers (default: 4)')
parser.add_argument('--epochs', default=30, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=128, type=int,metavar='N',
help='mini-batch size (default: 256), this is the total '
'batch size of all GPUs on the current node when '
'using Data Parallel or Distributed Data Parallel')
parser.add_argument('--lr', '--learning-rate', default=0.001, type=float,
metavar='LR', help='initial learning rate', dest='lr')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
help='momentum')
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
metavar='W', help='weight decay (default: 1e-4)',
dest='weight_decay')
parser.add_argument('-p', '--print-freq', default=10, type=int,
metavar='N', help='print frequency (default: 10)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('-e', '--evaluate', default=False, dest='evaluate', action='store_true',
help='evaluate model on validation set')
parser.add_argument('--pretrained', default='./trained_weights/fan_bestacc_PY37.pth', dest='pretrained', action='store_true',
help='use pre-trained model')
parser.add_argument('--world-size', default=-1, type=int,
help='number of nodes for distributed training')
parser.add_argument('--rank', default=-1, type=int,
help='node rank for distributed training')
parser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str,
help='url used to set up distributed training')
parser.add_argument('--dist-backend', default='nccl', type=str,
help='distributed backend')
parser.add_argument('--seed', default=None, type=int,
help='seed for initializing training. ')
parser.add_argument('--gpu', default=0, type=int,
help='GPU id to use.')
parser.add_argument('--multiprocessing-distributed', default=False, action='store_true',
help='Use multi-processing distributed training to launch '
'N processes per node, which has N GPUs. This is the '
'fastest way to use PyTorch for either single node or '
'multi node data parallel training')
parser.add_argument('--target_sample_rate', type=int, default=16000, help='Target sample rate for audio files')
parser.add_argument('--num_samples', type=int, default=160000, help='Number of samples to process')
parser.add_argument('--device', type=str, default='cuda', help='Device to use for training (e.g., "cuda" or "cpu")')
# parser.add_argument('--batch_size', type=int, default=256, help='Batch size for training/validation')
parser.add_argument('--model_path', type=str, default=None)
parser.add_argument('--backend', default='tensorrt', type=str, choices=['tengine_u8', 'tensorrt', 'nnie', 'ppl', 'snpe'])
parser.add_argument('--optim', type=str, default='adam')
parser.add_argument('--not-quant', action='store_true', default=False)
parser.add_argument('--deploy', action='store_true', default=False)
BackendMap = {'tensorrt': BackendType.Tensorrt,
'nnie': BackendType.NNIE,
'ppl': BackendType.PPLW8A16,
'snpe': BackendType.SNPE,
'vitis': BackendType.Vitis,
'tengine_u8': BackendType.Tengine_u8}
best_acc1 = 0
tb_idx = 0
prepare_custom_config_dict = dict(extra_qconfig_dict={
'w_observer': 'MinMaxObserver',
'a_observer': 'EMAMinMaxObserver',
'w_fakequantize': 'LearnableFakeQuantize',
'a_fakequantize': 'LearnableFakeQuantize',
'w_qscheme': {
'bit': 8,
'symmetry': True,
'per_channel': False,
'pot_scale': False
},
'a_qscheme': {
'bit': 8,
'symmetry': True,
'per_channel': False,
'pot_scale': False
}
},
extra_quantizer_dict={
'additional_node_name': 'output'
}
)
def main():
args = parser.parse_args()
args.quant = not args.not_quant
args.backend = BackendMap[args.backend]
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
if args.gpu is not None:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
if args.dist_url == "env://" and args.world_size == -1:
args.world_size = int(os.environ["WORLD_SIZE"])
args.distributed = args.world_size > 1 or args.multiprocessing_distributed
ngpus_per_node = torch.cuda.device_count()
if args.multiprocessing_distributed:
# Since we have ngpus_per_node processes per node, the total world_size
# needs to be adjusted accordingly
args.world_size = ngpus_per_node * args.world_size
# Use torch.multiprocessing.spawn to launch distributed processes: the
# main_worker process function
mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
else:
# Simply call main_worker function
main_worker(args.gpu, ngpus_per_node, args)
def main_worker(gpu, ngpus_per_node, args, tb_path='./logs/org'):
tb_writer = SummaryWriter(tb_path)
global best_acc1
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
world_size=args.world_size, rank=args.rank)
# create model
if args.pretrained is not None:
print("=> using pre-trained model '{}'".format(args.arch))
model = cnn.CNN()
model.load_state_dict(torch.load(args.pretrained))
else:
print("=> creating model '{}'".format(args.arch))
# model = cnn()
model = cnn.CNN() # CNN 是 cnn 模块中的一个类
# for internal cluster
if args.model_path is not None:
state_dict = torch.load(args.model_path)
print(f'load pretrained checkpoint from: {args.model_path}')
if 'model' in state_dict:
model.load_state_dict(state_dict['model'])
else:
state_dict = state_dict['state_dict']
for k, v in deepcopy(state_dict).items():
new_k = k.replace('module.', '') if 'module' in k else k
state_dict[new_k] = v
model.load_state_dict(state_dict)
# quantize model
if args.quant:
model = prepare_by_platform(model, args.backend, prepare_custom_config_dict)
if not torch.cuda.is_available():
print('using CPU, this will be slow')
elif args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
if args.optim == 'sgd':
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optim == 'adam':
optimizer = torch.optim.Adam(model.parameters(), args.lr,
betas=(0.9, 0.999), eps=1e-08,
weight_decay=args.weight_decay,
amsgrad=False)
# prepare dataset
train_loader, train_sampler, val_loader, cali_loader = prepare_dataloader(args)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
if args.gpu is None:
checkpoint = torch.load(args.resume)
else:
# Map model to be loaded to specified single gpu.
loc = 'cuda:{}'.format(args.gpu)
checkpoint = torch.load(args.resume, map_location=loc)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
if args.gpu is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpu)
state_dict = checkpoint['state_dict']
model_dict = model.state_dict()
if 'module.' in list(state_dict.keys())[0] and 'module.' not in list(model_dict.keys())[0]:
for k in list(state_dict.keys()):
state_dict[k[7:]] = state_dict.pop(k)
torch.save(checkpoint['state_dict'], './CNN_fan.pth')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {}), acc = {}"
.format(args.resume, checkpoint['epoch'], best_acc1))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
elif args.quant:
enable_calibration(model)
calibrate(cali_loader, model, args)
cudnn.benchmark = True
if args.quant:
enable_quantization(model)
if args.quant and args.deploy:
convert_deploy(model.eval(), args.backend, input_shape_dict={'data': [1, 1, args.audio_height, args.audio_width]},
output_path='.',
input_names=['input'],
output_names=['output'])
return
if args.evaluate:
if args.quant:
from mqbench.convert_deploy import convert_merge_bn
convert_merge_bn(model.eval())
validate(val_loader, model, criterion, args)
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
# tb_idx = train(train_loader, model, criterion, optimizer, epoch, args, tb_writer, tb_idx)
# evaluate on validation set
acc1 = validate(val_loader, model, criterion, args)
tb_writer.add_scalar('acc1_epoch', acc1, epoch)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if not args.multiprocessing_distributed or (args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0):
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer' : optimizer.state_dict(),
}, is_best)
# 处理音频数据
def prepare_dataloader(args):
# 设置音频数据的路径
# 计算Mel频谱图
MelSpectrogram = torchaudio.transforms.MelSpectrogram(
n_fft=400,
hop_length=160,
window_fn=torch.hamming_window
) # default
'''
train_paths = os.listdir(os.path.join(args.train_data, 'train'))
val_paths = os.listdir(os.path.join(args.val_data, 'val'))
'''
dataset_path = os.listdir(args.dataset)
DataSet = MachineDataset(
audio_dir=args.dataset,
paths=dataset_path,
transformation=MelSpectrogram,
target_sample_rate=args.target_sample_rate,
num_samples=args.num_samples,
device=args.gpu
)
# split dataset
train_size = int(0.7 * len(DataSet))
val_size = len(DataSet) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(DataSet, [train_size, val_size])
'''
# 加载训练数据集
train_dataset = MachineDataset(
audio_dir=os.path.join(args.train_data, 'train'),
paths=train_paths,
transformation=MelSpectrogram,
target_sample_rate=args.target_sample_rate,
num_samples=args.num_samples,
device=args.gpu
)
# 加载验证数据集
val_dataset = MachineDataset(
audio_dir=os.path.join(args.val_data, 'val'),
paths=val_paths,
transformation=MelSpectrogram,
target_sample_rate=args.target_sample_rate,
num_samples=args.num_samples,
device=args.gpu
)
'''
# 如果是分布式训练
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
else:
train_sampler = None
# 定义训练数据加载器
train_loader = DataLoader(
train_dataset, batch_size=64, shuffle=(train_sampler is None),
num_workers=args.workers, pin_memory=False, sampler=train_sampler
)
# 定义校准数据集和加载器
cali_batch_size = 10
cali_batch = 10
cali_dataset = torch.utils.data.Subset(train_dataset, indices=torch.arange(cali_batch_size * cali_batch))
cali_loader = DataLoader(cali_dataset, batch_size=cali_batch_size, shuffle=False,
num_workers=args.workers, pin_memory=False)
# 定义验证数据加载器
val_loader = DataLoader(
val_dataset, batch_size=32, shuffle=False,
num_workers=args.workers, pin_memory=False
)
return train_loader, train_sampler, val_loader, cali_loader
'''
# 图像
def calibrate(cali_loader, model, args):
model.eval()
print("Start calibration ...")
print("Calibrate images number = ", len(cali_loader.dataset))
with torch.no_grad():
for i, (images, target) in enumerate(cali_loader):
if args.gpu is not None:
images = images.cuda(args.gpu, non_blocking=True)
output = model(images)
print("Calibration ==> ", i+1)
print("End calibration.")
return
'''
# 音频处理
def calibrate(cali_loader, model, args):
model.eval()
print("Start calibration ...")
print("Calibrate audio files number = ", len(cali_loader.dataset))
with torch.no_grad():
for i, (audio, target) in enumerate(cali_loader):
if args.gpu is not None:
audio = audio.cuda(args.gpu, non_blocking=True)
output = model(audio) # 然后将变换后的数据输入模型
print("Calibration ==> ", i + 1)
print("End calibration.")
return
def train(train_loader, model, criterion, optimizer, epoch, args, tb_writer, tb_idx):
batch_time = AverageMeter('Time', ':6.3f')
data_time = AverageMeter('Data', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
top2 = AverageMeter('Acc@2', ':6.2f')
progress = ProgressMeter(
len(train_loader),
[batch_time, data_time, losses, top1, top2],
prefix="Epoch: [{}]".format(epoch))
# switch to train mode
model.train()
end = time.time()
for i, (images, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
if args.gpu is not None:
images = images.cuda(args.gpu, non_blocking=True)
if torch.cuda.is_available():
target = target.cuda(args.gpu, non_blocking=True)
# compute output
output = model(images)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc2 = accuracy(output, target, topk=(1, 2))
tb_writer.add_scalar('loss', loss, tb_idx)
tb_writer.add_scalar('acc1', acc1, tb_idx)
tb_writer.add_scalar('acc2', acc2, tb_idx)
losses.update(loss.item(), images.size(0))
top1.update(acc1[0], images.size(0))
top2.update(acc2[0], images.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
progress.display(i)
tb_idx += 1
return tb_idx
def validate(val_loader, model, criterion, args):
batch_time = AverageMeter('Time', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
top2 = AverageMeter('Acc@2', ':6.2f')
progress = ProgressMeter(
len(val_loader),
[batch_time, losses, top1, top2],
prefix='Test: ')
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
for i, (audio, target) in enumerate(val_loader):
if args.gpu is not None:
audio = audio.cuda(args.gpu, non_blocking=True)
if torch.cuda.is_available():
target = target.cuda(args.gpu, non_blocking=True)
# compute output
output = model(audio)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc2 = accuracy(output, target, topk=(1, 2))
losses.update(loss.item(), audio.size(0))
top1.update(acc1[0], audio.size(0))
top2.update(acc2[0], audio.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
progress.display(i)
# TODO: this should also be done with the ProgressMeter
print(' * Acc@1 {top1.avg:.3f} Acc@2 {top2.avg:.3f}'
.format(top1=top1, top2=top2))
return top1.avg
def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
torch.save(state, filename)
if is_best:
shutil.copyfile(filename, 'model_best.pth.tar')
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self, name, fmt=':f'):
self.name = name
self.fmt = fmt
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def __str__(self):
fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
return fmtstr.format(**self.__dict__)
class ProgressMeter(object):
def __init__(self, num_batches, meters, prefix=""):
self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
self.meters = meters
self.prefix = prefix
def display(self, batch):
entries = [self.prefix + self.batch_fmtstr.format(batch)]
entries += [str(meter) for meter in self.meters]
print('\t'.join(entries))
def _get_batch_fmtstr(self, num_batches):
num_digits = len(str(num_batches // 1))
fmt = '{:' + str(num_digits) + 'd}'
return '[' + fmt + '/' + fmt.format(num_batches) + ']'
def adjust_learning_rate(optimizer, epoch, args):
"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
lr = args.lr * (0.1 ** (epoch // 30))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def accuracy(output, target, topk=(1,)):
"""Computes the accuracy over the k top predictions for the specified values of k"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(100.0 / batch_size))
return res
if __name__ == '__main__':
main()
想和大神交朋友或想软件开发兼职接项目,请通过手机端搜小#程#序: "黄页小艺” 。
这段 Python 代码实现了一个基于 PyTorch 的音频故障检测模型的训练和评估流程,主要功能包括以下几个方面:
需软件开发兼职接项目,请通过手机端搜小#程#序: "黄页小艺” 。
一、命令行参数解析
使用argparse模块定义了一系列命令行参数,包括数据集路径、模型架构、训练参数(如批量大小、学习率、迭代次数等)、设备选择、量化和部署相关参数等。
二、主函数和工作进程函数
main函数:
处理随机种子设置,以确保可重复性。
根据是否使用分布式训练,调用mp.spawn启动多个进程进行分布式训练,或者直接调用main_worker进行单进程训练。
main_worker函数:
根据参数设置设备,处理分布式训练的初始化。
创建模型,如果有预训练模型则加载预训练权重。
根据参数进行模型量化和分布式数据并行设置。
定义损失函数和优化器。
准备数据集和数据加载器。
根据参数决定是否从 checkpoint 恢复训练、进行校准、评估或进行完整的训练和验证循环。
三、数据准备
prepare_dataloader函数:
读取音频数据的路径列表,并创建MachineDataset对象,该对象对音频数据进行预处理,包括计算 Mel 频谱图等操作。
将数据集划分为训练集和验证集。
根据是否为分布式训练创建相应的数据加载器和采样器。
四、校准和训练函数
calibrate函数:
对模型进行校准,在评估模式下遍历校准数据集,将数据输入模型但不计算梯度,用于量化过程中的参数调整。
train函数:
进行模型的训练循环,包括计算损失、准确率,进行梯度清零、反向传播和优化器更新等操作,并记录训练过程中的时间、损失和准确率等指标。
五、验证和保存 checkpoint 函数
validate函数:
在评估模式下对验证集进行评估,计算损失和准确率,并打印结果。
save_checkpoint函数:
保存模型的 checkpoint,如果当前模型是最佳模型,则复制为model_best.pth.tar。
六、辅助类和函数
AverageMeter类:用于跟踪和计算平均值的辅助类。
ProgressMeter类:用于显示训练进度的辅助类。
adjust_learning_rate函数:根据 epoch 调整学习率。
accuracy函数:计算模型在给定 top-k 值下的准确率。
想和大神交朋友或想软件开发兼职接项目,请通过手机端搜小#程#序: "黄页小艺” 。
总体来说,这段代码实现了一个音频故障检测模型的训练、评估和量化流程,支持分布式训练和多种优化器选择,以及模型的保存和加载。。.
魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐


所有评论(0)