SSD学习记录
前言
通过对目标确定四个参数，分别是目标中心点的x轴、y轴坐标，目标的框的高、宽，来确定目标的位置。
SSD首先会将图片调整成300x300的大小，为了防止失真，会在图像边缘加上灰条。
然后将图片分成1x1、 3x3、 5x5、 10x10、 19x19、 38x38，6种不同的大小的网格。网格分割越小，越有利于检测较小的目标，且每个网格中心负责每个目标。
总之，SSD是将一张图片划分成不同的网格，不某一个目标的中心点落在这个区域，这个物体就由这个网格来确定。
环境配制及相关资料
转载：https://blog.csdn.net/weixin_44791964/article/details/106037141
转载：https://www.bilibili.com/video/BV1A7411976Z?p=1
转载：https://blog.csdn.net/weixin_44791964/article/details/104981486
网络结构
网络基础介绍
- 将输入的图片resize成300x300x3的图片
- 通过VGG的主干网络模型，进行特征层的提取。
例如：38x38x512的特征层就分割成38x38的网格。 - 分割成网格之后，就有利于生成先验框。
例如：38x38x512的特征层对应4个先验框。该层就有38x38x4=5776 个先验框，判断这5776个先验框中是否包含我们所检测的目标。总共有8732个先验框。 - 通过非极大抑制，来找到我们所需要的先验框，并且标出类型。
先验框介绍
# 先验框主干部分
Config = {
'num_classes': 21,
'min_dim': 300,
'feature_maps': {
'vgg' : [38, 19, 10, 5, 3, 1],
'mobilenet' : [19, 10, 5, 3, 2, 1],
},
'min_sizes': [30, 60, 111, 162, 213, 264],
'max_sizes': [60, 111, 162, 213, 264, 315],
'aspect_ratios': {
'vgg' : [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
'mobilenet' : [[2, 3], [2, 3], [2, 3], [2, 3], [2, 3], [2, 3]]
},
'variance': [0.1, 0.2],
'clip': True,
}
mean = []
for k, f in enumerate(Config["feature_maps"]["vgg"]):
x,y = np.meshgrid(np.arange(f),np.arange(f))
x = x.reshape(-1)
y = y.reshape(-1)
for i, j in zip(y,x):
# print(x,y)
# 300/8
f_k = Config["feature_maps"]["vgg"][k]
# 计算网格的中心
cx = (j + 0.5) / f_k
cy = (i + 0.5) / f_k
# 求短边的正方形
s_k = Config["min_sizes"][k]/Config["min_dim"]
mean += [cx, cy, s_k, s_k]
# 求长边的正方形
s_k_prime = sqrt(s_k * (Config["max_sizes"][k]/Config["min_dim"]))
mean += [cx, cy, s_k_prime, s_k_prime]
# 获得长方形
for ar in Config["aspect_ratios"]["vgg"][k]:
mean += [cx, cy, s_k*sqrt(ar), s_k/sqrt(ar)]
mean += [cx, cy, s_k/sqrt(ar), s_k*sqrt(ar)]
mean = np.clip(mean,0,1)
mean = np.reshape(mean,[-1,4])*Config["min_dim"]
# 先验框打印
linx = np.linspace(0.5 * Config["min_dim"]/Config["feature_maps"]["vgg"][4], Config["min_dim"] - 0.5 * Config["min_dim"]/Config["feature_maps"]["vgg"][4],
Config["feature_maps"]["vgg"][4])
liny = np.linspace(0.5 * Config["min_dim"]/Config["feature_maps"]["vgg"][4], Config["min_dim"] - 0.5 * Config["min_dim"]/Config["feature_maps"]["vgg"][4],
Config["feature_maps"]["vgg"][4])
print("linx:",linx)
print("liny:",liny)
centers_x, centers_y = np.meshgrid(linx, liny)
fig = plt.figure()
ax = fig.add_subplot(111)
plt.ylim(-100,500)
plt.xlim(-100,500)
plt.scatter(centers_x,centers_y)
step_start = 8708
step_end = 8712
# step_start = 8728
# step_end = 8732
box_widths = mean[step_start:step_end,2]
box_heights = mean[step_start:step_end,3]
prior_boxes = np.zeros_like(mean[step_start:step_end,:])
prior_boxes[:,0] = mean[step_start:step_end,0]
prior_boxes[:,1] = mean[step_start:step_end,1]
prior_boxes[:,0] = mean[step_start:step_end,0]
prior_boxes[:,1] = mean[step_start:step_end,1]
# 获得先验框的左上角和右下角
prior_boxes[:, 0] -= box_widths/2
prior_boxes[:, 1] -= box_heights/2
prior_boxes[:, 2] += box_widths/2
prior_boxes[:, 3] += box_heights/2
rect1 = plt.Rectangle([prior_boxes[0, 0],prior_boxes[0, 1]],box_widths[0],box_heights[0],color="r",fill=False)
rect2 = plt.Rectangle([prior_boxes[1, 0],prior_boxes[1, 1]],box_widths[1],box_heights[1],color="r",fill=False)
rect3 = plt.Rectangle([prior_boxes[2, 0],prior_boxes[2, 1]],box_widths[2],box_heights[2],color="r",fill=False)
rect4 = plt.Rectangle([prior_boxes[3, 0],prior_boxes[3, 1]],box_widths[3],box_heights[3],color="r",fill=False)
ax.add_patch(rect1)
ax.add_patch(rect2)
ax.add_patch(rect3)
ax.add_patch(rect4)
plt.show()
print(np.shape(mean))
主干网络介绍
SSD模型的主干网络是经典的VGG网络结构。
# VGG部分
base = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
512, 512, 512]
def vgg(i):
layers = []
#i代表输入的通道数
in_channels = i
for v in base:
# M表示maxpool,C表示当前的in_channels
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
elif v == 'C':
layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
pool5 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6) # 起膨胀作用，相当于VGG中的全连接层。
conv7 = nn.Conv2d(1024, 1024, kernel_size=1)
layers += [pool5, conv6,
nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]
return layers
ceil_mode功能：
False时，不会对单边最大池化
True时，保留单边部分。
def add_extras(i, batch_norm=False):
# 添加到VGG的额外层用于功能缩放
layers = []
in_channels = i
# 先进行压缩，再进行通道调整
# Block 6
# 19,19,1024 -> 10,10,512
layers += [nn.Conv2d(in_channels, 256, kernel_size=1, stride=1)]
layers += [nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)]
# Block 7
# 10,10,512 -> 5,5,256
layers += [nn.Conv2d(512, 128, kernel_size=1, stride=1)]
layers += [nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)]
# Block 8
# 5,5,256 -> 3,3,256
layers += [nn.Conv2d(256, 128, kernel_size=1, stride=1)]
layers += [nn.Conv2d(128, 256, kernel_size=3, stride=1)]
# Block 9
# 3,3,256 -> 1,1,256
layers += [nn.Conv2d(256, 128, kernel_size=1, stride=1)]
layers += [nn.Conv2d(128, 256, kernel_size=3, stride=1)]
return layers
SSD模型构建
# 判断先验框中是否真实的包含物体，将先验框进行调整成预测框
class SSD(nn.Module):
def __init__(self, phase, base, extras, head, num_classes):
super(SSD, self).__init__()
self.phase = phase
self.num_classes = num_classes
self.cfg = Config
self.vgg = nn.ModuleList(base)
self.L2Norm = L2Norm(512, 20)
self.extras = nn.ModuleList(extras)
self.priorbox = PriorBox(self.cfg)
with torch.no_grad():
self.priors = Variable(self.priorbox.forward())
self.loc = nn.ModuleList(head[0])
self.conf = nn.ModuleList(head[1])
if phase == 'test':
self.softmax = nn.Softmax(dim=-1)
self.detect = Detect(num_classes, 0, 200, 0.01, 0.45)
def forward(self, x):
sources = list() # 用于回归预测与分类预测的连接
loc = list()
conf = list()
# 获得conv4_3的内容
for k in range(23):
x = self.vgg[k](x)
s = self.L2Norm(x) # 标准化
sources.append(s)
# 获得fc7的内容
for k in range(23, len(self.vgg)):
x = self.vgg[k](x)
sources.append(x)
# 获得后面的内容
for k, v in enumerate(self.extras):
x = F.relu(v(x), inplace=True)
if k % 2 == 1:
sources.append(x)
# 添加回归层和分类层。
# contiguous对通道数进行翻转，方便读取通道数。
for (x, l, c) in zip(sources, self.loc, self.conf):
loc.append(l(x).permute(0, 2, 3, 1).contiguous())
conf.append(c(x).permute(0, 2, 3, 1).contiguous())
# 进行resize
loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
if self.phase == "test":
# loc会resize到batch_size,num_anchors,4
# conf会resize到batch_size,num_anchors,
output = self.detect(
loc.view(loc.size(0), -1, 4), # 每个先验框的调整参数
self.softmax(conf.view(conf.size(0), -1,
self.num_classes)), # 每个先验框的种类
self.priors
)
else:
output = (
loc.view(loc.size(0), -1, 4),
conf.view(conf.size(0), -1, self.num_classes),
self.priors
)
return output
mbox = [4, 6, 6, 6, 4, 4]
# 分类预测与回归预测
def get_ssd(phase,num_classes):
vgg, extra_layers = add_vgg(3), add_extras(1024)
loc_layers = []
conf_layers = []
vgg_source = [21, -2]
for k, v in enumerate(vgg_source):
# 回归预测
loc_layers += [nn.Conv2d(vgg[v].out_channels,
mbox[k] * 4, kernel_size=3, padding=1)]
# 分类预测
conf_layers += [nn.Conv2d(vgg[v].out_channels,
mbox[k] * num_classes, kernel_size=3, padding=1)]
for k, v in enumerate(extra_layers[1::2], 2):
loc_layers += [nn.Conv2d(v.out_channels, mbox[k]
* 4, kernel_size=3, padding=1)]
conf_layers += [nn.Conv2d(v.out_channels, mbox[k]
* num_classes, kernel_size=3, padding=1)]
SSD_MODEL = SSD(phase, vgg, extra_layers, (loc_layers, conf_layers), num_classes)
return SSD_MODEL
获得预测框
class Detect(Function):
def __init__(self, num_classes, bkg_label, top_k, conf_thresh, nms_thresh):
self.num_classes = num_classes
self.background_label = bkg_label
self.top_k = top_k
self.nms_thresh = nms_thresh
if nms_thresh <= 0:
raise ValueError('nms_threshold must be non negative.')
self.conf_thresh = conf_thresh
self.variance = Config['variance']
def forward(self, loc_data, conf_data, prior_data):
loc_data = loc_data.cpu()
conf_data = conf_data.cpu()
num = loc_data.size(0) # batch_size
num_priors = prior_data.size(0) # 先验框数量9732
output = torch.zeros(num, self.num_classes, self.top_k, 5) # 空阵列，用于存放输出
conf_preds = conf_data.view(num, num_priors,
self.num_classes).transpose(2, 1)
# 对每一张图片进行处理
for i in range(num):
# 对先验框解码获得预测框
decoded_boxes = decode(loc_data[i], prior_data, self.variance)
conf_scores = conf_preds[i].clone()
for cl in range(1, self.num_classes):
# 对每一类进行非极大抑制
c_mask = conf_scores[cl].gt(self.conf_thresh)
scores = conf_scores[cl][c_mask]
if scores.size(0) == 0:
continue
l_mask = c_mask.unsqueeze(1).expand_as(decoded_boxes)
boxes = decoded_boxes[l_mask].view(-1, 4)
# 进行非极大抑制
ids, count = nms(boxes, scores, self.nms_thresh, self.top_k)
output[i, cl, :count] = \
torch.cat((scores[ids[:count]].unsqueeze(1),
boxes[ids[:count]]), 1)
flt = output.contiguous().view(num, -1, 5)
_, idx = flt[:, :, 0].sort(1, descending=True)
_, rank = idx.sort(1)
flt[(rank < self.top_k).unsqueeze(-1).expand_as(flt)].fill_(0)
return output
# 确定先验框的位置
def decode(loc, priors, variances):
boxes = torch.cat((
priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], # 调整后的中心
priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1) # 调整后的高宽
boxes[:, :2] -= boxes[:, 2:] / 2
boxes[:, 2:] += boxes[:, :2]
return boxes
预测过程
MEANS = (104, 117, 123)
class SSD(object):
_defaults = {
"model_path" : 'model_data/ssd_weights.pth', # 权重
"classes_path" : 'model_data/voc_classes.txt', # 分类
"input_shape" : (300, 300, 3), # 输入图片的大小
"confidence" : 0.5, # 可以容忍的置信度
"nms_iou" : 0.45,
"cuda" : True,
"backbone" : "vgg",
"letterbox_image" : False,
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
self.class_names = self._get_class()
self.generate()
def _get_class(self):
classes_path = os.path.expanduser(self.classes_path)
with open(classes_path) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
def generate(self):
self.num_classes = len(self.class_names) + 1
model = ssd.get_ssd("test", self.num_classes, self.backbone, self.confidence, self.nms_iou)
print('Loading weights into state dict...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(torch.load(self.model_path, map_location=device))
self.net = model.eval()
if self.cuda:
self.net = torch.nn.DataParallel(self.net)
cudnn.benchmark = True
self.net = self.net.cuda()
print('{} model, anchors, and classes loaded.'.format(self.model_path))
# 画框设置不同的颜色
hsv_tuples = [(x / len(self.class_names), 1., 1.)
for x in range(len(self.class_names))]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
self.colors))
# 检测图片
def detect_image(self, image):
image = image.convert('RGB')
image_shape = np.array(np.shape(image)[0:2])
# 增加灰条，防止失真
if self.letterbox_image:
crop_img = np.array(letterbox_image(image, (self.input_shape[1],self.input_shape[0])))
else:
crop_img = image.resize((self.input_shape[1],self.input_shape[0]), Image.BICUBIC)
# 图片预处理，归一化
photo = np.array(crop_img,dtype = np.float64)
with torch.no_grad():
photo = torch.from_numpy(np.expand_dims(np.transpose(photo - MEANS, (2, 0, 1)), 0)).type(torch.FloatTensor)
if self.cuda:
photo = photo.cuda()
preds = self.net(photo)
top_conf = []
top_label = []
top_bboxes = []
# 置信度的筛选
for i in range(preds.size(1)):
j = 0
while preds[0, i, j, 0] >= self.confidence:
score = preds[0, i, j, 0]
label_name = self.class_names[i-1]
pt = (preds[0, i, j, 1:]).detach().numpy()
coords = [pt[0], pt[1], pt[2], pt[3]]
top_conf.append(score)
top_label.append(label_name)
top_bboxes.append(coords)
j = j + 1
# 如果不存在满足门限的预测框，直接返回原图
if len(top_conf)<=0:
return image
top_conf = np.array(top_conf)
top_label = np.array(top_label)
top_bboxes = np.array(top_bboxes)
top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:,0], -1),np.expand_dims(top_bboxes[:,1], -1),np.expand_dims(top_bboxes[:,2], -1),np.expand_dims(top_bboxes[:,3], -1)
# 去掉灰条
if self.letterbox_image:
boxes = ssd_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.input_shape[0],self.input_shape[1]]),image_shape)
else:
top_xmin = top_xmin * image_shape[1]
top_ymin = top_ymin * image_shape[0]
top_xmax = top_xmax * image_shape[1]
top_ymax = top_ymax * image_shape[0]
boxes = np.concatenate([top_ymin,top_xmin,top_ymax,top_xmax], axis=-1)
font = ImageFont.truetype(font='model_data/simhei.ttf',size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32'))
thickness = max((np.shape(image)[0] + np.shape(image)[1]) // self.input_shape[0], 1)
for i, c in enumerate(top_label):
predicted_class = c
score = top_conf[i]
top, left, bottom, right = boxes[i]
top = top - 5
left = left - 5
bottom = bottom + 5
right = right + 5
top = max(0, np.floor(top + 0.5).astype('int32'))
left = max(0, np.floor(left + 0.5).astype('int32'))
bottom = min(np.shape(image)[0], np.floor(bottom + 0.5).astype('int32'))
right = min(np.shape(image)[1], np.floor(right + 0.5).astype('int32'))
# 画框框
label = '{} {:.2f}'.format(predicted_class, score)
draw = ImageDraw.Draw(image)
label_size = draw.textsize(label, font)
label = label.encode('utf-8')
print(label, top, left, bottom, right)
if top - label_size[1] >= 0:
text_origin = np.array([left, top - label_size[1]])
else:
text_origin = np.array([left, top + 1])
for i in range(thickness):
draw.rectangle(
[left + i, top + i, right - i, bottom - i],
outline=self.colors[self.class_names.index(predicted_class)])
draw.rectangle(
[tuple(text_origin), tuple(text_origin + label_size)],
fill=self.colors[self.class_names.index(predicted_class)])
draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
del draw
return image
训练部分
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
def fit_one_epoch(net,criterion,epoch,epoch_size,epoch_size_val,gen,genval,Epoch,cuda):
loc_loss = 0
conf_loss = 0
loc_loss_val = 0
conf_loss_val = 0
net.train()
print('Start Train')
with tqdm(total=epoch_size,desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar:
for iteration, batch in enumerate(gen):
if iteration >= epoch_size:
break
images, targets = batch[0], batch[1]
with torch.no_grad():
if cuda:
images = torch.from_numpy(images).type(torch.FloatTensor).cuda()
targets = [torch.from_numpy(ann).type(torch.FloatTensor).cuda() for ann in targets]
else:
images = torch.from_numpy(images).type(torch.FloatTensor)
targets = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets]
# 前向传播
out = net(images)
# 清零梯度
optimizer.zero_grad()
# 计算损失
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
# 反向传播
loss.backward()
optimizer.step()
loc_loss += loss_l.item()
conf_loss += loss_c.item()
pbar.set_postfix(**{'loc_loss' : loc_loss / (iteration + 1),
'conf_loss' : conf_loss / (iteration + 1),
'lr' : get_lr(optimizer)})
pbar.update(1)
net.eval()
print('Start Validation')
with tqdm(total=epoch_size_val, desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar:
for iteration, batch in enumerate(genval):
if iteration >= epoch_size_val:
break
images, targets = batch[0], batch[1]
with torch.no_grad():
if cuda:
images = torch.from_numpy(images).type(torch.FloatTensor).cuda()
targets = [torch.from_numpy(ann).type(torch.FloatTensor).cuda() for ann in targets]
else:
images = torch.from_numpy(images).type(torch.FloatTensor)
targets = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets]
out = net(images)
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loc_loss_val += loss_l.item()
conf_loss_val += loss_c.item()
pbar.set_postfix(**{'loc_loss' : loc_loss_val / (iteration + 1),
'conf_loss' : conf_loss_val / (iteration + 1),
'lr' : get_lr(optimizer)})
pbar.update(1)
total_loss = loc_loss + conf_loss
val_loss = loc_loss_val + conf_loss_val
loss_history.append_loss(total_loss/(epoch_size+1), val_loss/(epoch_size_val+1))
print('Finish Validation')
print('Epoch:'+ str(epoch+1) + '/' + str(Epoch))
print('Total Loss: %.4f || Val Loss: %.4f ' % (total_loss/(epoch_size+1),val_loss/(epoch_size_val+1)))
print('Saving state, iter:', str(epoch+1))
torch.save(model.state_dict(), 'logs/Epoch%d-Total_Loss%.4f-Val_Loss%.4f.pth'%((epoch+1),total_loss/(epoch_size+1),val_loss/(epoch_size_val+1)))
return val_loss/(epoch_size_val+1)
# 检测精度mAP和pr曲线计算参考视频
if __name__ == "__main__":
# 是否使用Cuda
# 没有GPU可以设置成False
Cuda = True
# 与视频中不同、新添加了主干网络的选择
# 分别实现了基于mobilenetv2和vgg的ssd
# 可通过修改backbone变量进行主干网络的选择
# vgg或者mobilenet
backbone = "vgg"
model = get_ssd("train", Config["num_classes"], backbone)
weights_init(model)
# 权值文件请看README，百度网盘下载
model_path = "model_data/ssd_weights.pth"
print('Loading weights into state dict...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dict = model.state_dict()
pretrained_dict = torch.load(model_path, map_location=device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
print('Finished!')
annotation_path = '2007_train.txt'
# 验证集的划分在train.py代码里面进行
# 2007_test.txt和2007_val.txt里面没有内容是正常的。训练不会使用到。
# 当前划分方式下，验证集和训练集的比例为1:9
val_split = 0.1
with open(annotation_path) as f:
lines = f.readlines()
np.random.seed(10101)
np.random.shuffle(lines)
np.random.seed(None)
num_val = int(len(lines)*val_split)
num_train = len(lines) - num_val
criterion = MultiBoxLoss(Config['num_classes'], 0.5, True, 0, True, 3, 0.5, False, Cuda)
loss_history = LossHistory("logs/")
net = model.train()
if Cuda:
net = torch.nn.DataParallel(model)
cudnn.benchmark = True
net = net.cuda()
# 主干特征提取网络特征通用，冻结训练可以加快训练速度
# 也可以在训练初期防止权值被破坏。
# Init_Epoch为起始世代
# Freeze_Epoch为冻结训练的世代
# Unfreeze_Epoch总训练世代
# 提示OOM或者显存不足请调小Batch_size
if True:
lr = 5e-4
Batch_size = 32
Init_Epoch = 0
Freeze_Epoch = 50
optimizer = optim.Adam(net.parameters(), lr=lr)
lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3, verbose=True)
train_dataset = SSDDataset(lines[:num_train], (Config["min_dim"], Config["min_dim"]), True)
val_dataset = SSDDataset(lines[num_train:], (Config["min_dim"], Config["min_dim"]), False)
gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=ssd_dataset_collate)
gen_val = DataLoader(val_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=ssd_dataset_collate)
if backbone == "vgg":
for param in model.vgg.parameters():
param.requires_grad = False
else:
for param in model.mobilenet.parameters():
param.requires_grad = False
epoch_size = num_train // Batch_size
epoch_size_val = num_val // Batch_size
if epoch_size == 0 or epoch_size_val == 0:
raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
for epoch in range(Init_Epoch,Freeze_Epoch):
val_loss = fit_one_epoch(net,criterion,epoch,epoch_size,epoch_size_val,gen,gen_val,Freeze_Epoch,Cuda)
lr_scheduler.step(val_loss)
if True:
lr = 1e-4
Batch_size = 16
Freeze_Epoch = 50
Unfreeze_Epoch = 100
optimizer = optim.Adam(net.parameters(), lr=lr)
lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3, verbose=True)
train_dataset = SSDDataset(lines[:num_train], (Config["min_dim"], Config["min_dim"]), True)
val_dataset = SSDDataset(lines[num_train:], (Config["min_dim"], Config["min_dim"]), False)
gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=ssd_dataset_collate)
gen_val = DataLoader(val_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=ssd_dataset_collate)
if backbone == "vgg":
for param in model.vgg.parameters():
param.requires_grad = True
else:
for param in model.mobilenet.parameters():
param.requires_grad = True
epoch_size = num_train // Batch_size
epoch_size_val = num_val // Batch_size
if epoch_size == 0 or epoch_size_val == 0:
raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
for epoch in range(Freeze_Epoch,Unfreeze_Epoch):
val_loss = fit_one_epoch(net,criterion,epoch,epoch_size,epoch_size_val,gen,gen_val,Unfreeze_Epoch,Cuda)
lr_scheduler.step(val_loss)
484
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
