utils.general详解1

class WorkingDirectory(contextlib.ContextDecorator): # 改变函数的工作路径
    # Usage: @WorkingDirectory(dir) decorator or 'with WorkingDirectory(dir):' context manager
    def __init__(self, new_dir):
        self.dir = new_dir  # new dir
        self.cwd = Path.cwd().resolve()  # current dir

    def __enter__(self):
        os.chdir(self.dir)

    def __exit__(self, exc_type, exc_val, exc_tb):
        os.chdir(self.cwd)

def try_except(func):
    # try-except function. Usage: @try_except decorator
    def handler(*args, **kwargs):
        try:
            func(*args, **kwargs)
        except Exception as e:
            print(e)

    return handler

WorkingDirectory类用来控制函数的运行路径如下例子所示

from posixpath import normpath
from cv2 import trace
from utils.general import check_git_status,WorkingDirectory, try_except
from pathlib import Path
import os
FILE = Path(__file__).resolve()
ROOT = FILE.parents[1]  # YOLOv5 root directory

@try_except
@WorkingDirectory(r'E:\thesis')
def func():
    print('func')
    print(os.getcwd())

func()

WorkingDirectory作为装饰器,可以传入一个路径参数用来指定当前函数的工作路径,不指定的话就在当前路径,try_except也是装饰器,用来接收异常,不让程序在该函数出现异常时停止。

yolov7train.py 是使用 YOLOv7 算法进行目标检测的训练脚本。下面对 yolov7train.py 的主要代码进行简单的解释: 1. 导入相关库 ```python import argparse import yaml import time import torch from torch.utils.data import DataLoader from torchvision import datasets from models.yolov7 import Model from utils.datasets import ImageFolder from utils.general import ( check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box, strip_optimizer, set_logging) from utils.torch_utils import ( select_device, time_synchronized, load_classifier, model_info) ``` 这里导入了 argparse 用于解析命令行参数,yaml 用于解析配置文件,time 用于记录时间,torch 用于神经网络训练,DataLoader 用于读取数据集,datasets 和 ImageFolder 用于加载数据集,Model 用于定义 YOLOv7 模型,各种工具函数用于辅助训练。 2. 定义命令行参数 ```python parser = argparse.ArgumentParser() parser.add_argument('--data', type=str, default='data.yaml', help='dataset.yaml path') parser.add_argument('--hyp', type=str, default='hyp.yaml', help='hyperparameters path') parser.add_argument('--epochs', type=int, default=300) parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs') parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes') parser.add_argument('--rect', action='store_true', help='rectangular training') parser.add_argument('--resume', nargs='?', const='yolov7.pt', default=False, help='resume most recent training') parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') parser.add_argument('--notest', action='store_true', help='only test final epoch') parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters') parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') opt = parser.parse_args() ``` 这里定义了许多命令行参数,包括数据集路径、超参数路径、训练轮数、批量大小、图片大小、是否使用矩形训练、是否从最近的检查点恢复训练、是否只保存最终的检查点、是否只测试最终的模型、是否进行超参数进化、gsutil 存储桶等。 3. 加载数据集 ```python with open(opt.data) as f: data_dict = yaml.load(f, Loader=yaml.FullLoader) train_path = data_dict['train'] test_path = data_dict['test'] num_classes = data_dict['nc'] names = data_dict['names'] train_dataset = ImageFolder(train_path, img_size=opt.img_size[0], rect=opt.rect) test_dataset = ImageFolder(test_path, img_size=opt.img_size[1], rect=True) batch_size = opt.batch_size train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True, collate_fn=train_dataset.collate_fn) test_dataloader = DataLoader(test_dataset, batch_size=batch_size * 2, num_workers=8, pin_memory=True, collate_fn=test_dataset.collate_fn) ``` 这里读取了数据集的配置文件,包括训练集、测试集、类别数和类别名称等信息。然后使用 ImageFolder 加载数据集,设置图片大小和是否使用矩形训练。最后使用 DataLoader 加载数据集,并设置批量大小、是否 shuffle、是否使用 pin_memory 等参数。 4. 定义 YOLOv7 模型 ```python model = Model(opt.hyp, num_classes, opt.img_size) model.nc = num_classes device = select_device(opt.device, batch_size=batch_size) model.to(device).train() criterion = model.loss optimizer = torch.optim.SGD(model.parameters(), lr=hyp['lr0'], momentum=hyp['momentum'], weight_decay=hyp['weight_decay']) scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=1, T_mult=2) start_epoch = 0 best_fitness = 0.0 ``` 这里使用 Model 类定义了 YOLOv7 模型,并将其放到指定设备上进行训练。使用交叉熵损失函数作为模型的损失函数,使用 SGD 优化器进行训练,并使用余弦退火学习率调整策略。定义了起始轮数、最佳精度等变量。 5. 开始训练 ```python for epoch in range(start_epoch, opt.epochs): model.train() mloss = torch.zeros(4).to(device) # mean losses for i, (imgs, targets, paths, _) in enumerate(train_dataloader): ni = i + len(train_dataloader) * epoch # number integrated batches (since train start) imgs = imgs.to(device) targets = targets.to(device) loss, _, _ = model(imgs, targets) loss.backward() optimizer.step() optimizer.zero_grad() mloss = (mloss * i + loss.detach().cpu()) / (i + 1) # update mean losses # Print batch results if ni % 20 == 0: print(f'Epoch {epoch}/{opt.epochs - 1}, Batch {i}/{len(train_dataloader) - 1}, lr={optimizer.param_groups[0]["lr"]:.6f}, loss={mloss[0]:.4f}') # Update scheduler scheduler.step() # Update Best fitness with torch.no_grad(): fitness = model_fitness(model) if fitness > best_fitness: best_fitness = fitness # Save checkpoint if (not opt.nosave) or (epoch == opt.epochs - 1): ckpt = { 'epoch': epoch, 'best_fitness': best_fitness, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() } torch.save(ckpt, f'checkpoints/yolov7_epoch{epoch}.pt') # Test if not opt.notest: t = time_synchronized() model.eval() for j, (imgs, targets, paths, shapes) in enumerate(test_dataloader): if j == 0: pred = model(imgs.to(device)) pred = non_max_suppression(pred, conf_thres=0.001, iou_thres=0.6) else: break t1 = time_synchronized() if isinstance(pred, int) or isinstance(pred, tuple): print(f'Epoch {epoch}/{opt.epochs - 1}, test_loss={mloss[0]:.4f}, test_mAP={0.0}') else: pred = pred[0].cpu() iou_thres = 0.5 niou = [iou_thres] * num_classes ap, p, r = ap_per_class(pred, targets, shapes, iou_thres=niou) mp, mr, map50, f1, _, _ = stats(ap, p, r, gt=targets) print(f'Epoch {epoch}/{opt.epochs - 1}, test_loss={mloss[0]:.4f}, test_mAP={map50:.2f} ({mr*100:.1f}/{mp*100:.1f})') # Plot images if epoch == 0 and j == 0: for i, det in enumerate(pred): # detections per image img = cv2.imread(paths[i]) # BGR img = plot_results(img, det, class_names=names) cv2.imwrite(f'runs/test{i}.jpg', img) if i == 3: break ``` 这里进行了多个 epoch 的训练。在每个 epoch 中,对于每个批量的数据,先将数据移动到指定设备上,然后计算模型的损失函数,并进行反向传播和梯度下降。在每个 epoch 结束时,更新学习率调整策略和最佳精度,保存当前的检查点。如果 opt.notest 为 False,则进行测试,并输出测试结果。最后,如果是第一个 epoch,则绘制部分图像用于可视化。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值