#---------------step1:Dataset-------------------
import torch
from YOLO_V1_DataSet import VOCDataSet
dataSet = VOCDataSet(imgs_dir="./VOC2007/Train/JPEGImages",
annotations_dir="./VOC2007/Train/Annotations",
ClassesFile="./VOC2007/Train/class.data")
from torch.utils.data import DataLoader
#dataLoader = DataLoader(dataSet,batch_size=64,shuffle=True,num_workers=0)
#---------------step2:Model-------------------
from YOLO_V1_Model import YOLO_V1
YOLO = YOLO_V1().cuda(device=0)
YOLO.initialize_weights()
#---------------step3:LossFunction-------------------
from YOLO_V1_LossFunction import YOLO_V1_Loss
loss_function = YOLO_V1_Loss().cuda(device=0)
#---------------step4:Optimizer-------------------
import torch.optim as optim
optimizer_Adam = optim.Adam(YOLO.parameters(), lr=1e-4, weight_decay=0.0005)
#使用余弦退火动态调整学习率
#lr_reduce_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer_Adam , T_max=20, eta_min=1e-4, last_epoch=-1)
#lr_reduce_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer=optimizer_Adam, T_0=2, T_mult=2)
#--------------step5:Tensorboard Feature Map------------
from tensorboardX import SummaryWriter
import torchvision.utils as vutils
import torch.nn as nn
def feature_map_visualize(img_data, writer):
# 网络加载
from YOLO_V1_Model import YOLO_V1
YOLO = YOLO_V1()
import torch
YOLO.load_state_dict(torch.load('./YOLO_V1_300.pth'))
# 类别与索引转换
IndexToClassName = {}
with open("./VOC2007/Train/class.data", "r") as f:
index = 0
for line in f:
IndexToClassName[index] = line
index = index + 1
def iou(box_one, box_two):
LX = max(box_one[0], box_two[0])
LY = max(box_one[1], box_two[1])
RX = min(box_one[2], box_two[2])
RY = min(box_one[3], box_two[3])
if LX >= RX or LY >= RY:
return 0
return (RX - LX) * (RY - LY) / ((box_one[2] - box_one[0]) *
(box_one[3] - box_one[1]) +
(box_two[2] - box_two[0]) *
(box_two[3] - box_two[1]))
import numpy as np
import torch
from YOLO_V1_Model import YOLO_V1
weight_file_name = "YOLO_V1_2400.pth"
model = YOLO_V1().cuda(device=1)
model.load_state_dict(torch.load(weight_file_name))
torch.save(model.state_dict(),
weight_file_name,
_use_new_zipfile_serialization=False)
#------step1: 导入网络------
from YOLO_V1_Model import YOLO_V1
YOLO = YOLO_V1().cuda()
#------step:2 读取权重文件------
import torch
weight_file_name = "./Model_Train/YOLO_V1_500.pth"
YOLO.load_state_dict(torch.load(weight_file_name)["model"])
YOLO.eval()
#------step:3 类别索引与类别名的映射------
class_file_name = "./VOC2007/Train/class.data"
class_index_Name = {}
classIndex = 0
with open(class_file_name, 'r') as f:
for line in f:
line = line.replace('\n', '')
class_index_Name[classIndex] = line # 根据类别名制作索引
classIndex = classIndex + 1
#------step:4 NMS算法处理输出结果------
def iou(box_one, box_two):
LX = max(box_one[0], box_two[0])
LY = max(box_one[1], box_two[1])
RX = min(box_one[2], box_two[2])
RY = min(box_one[3], box_two[3])
if LX >= RX or LY >= RY:
return 0
return (RX - LX) * (RY - LY) / ((box_one[2] - box_one[0]) *