def generate(self):
self.confidence = 0.01
#-------------------------------#
# 计算总的类的数量
#-------------------------------#
self.num_classes = len(self.class_names) + 1
#-------------------------------#
# 载入模型与权值
#-------------------------------#
model = get_ssd("test", self.num_classes, 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 generate(self):
# 计算总的种类
self.num_classes = len(self.class_names) + 1
# 载入模型，如果原来的模型里已经包括了模型结构则直接载入。
# 否则先构建模型再载入
model = ssd.get_ssd("test", self.num_classes)
self.net = model
# 标签文件中不能有空白行
model.load_state_dict(torch.load(self.Config["model_path"],
map_location=torch.device('cpu')),
strict=False)
self.net = torch.nn.DataParallel(self.net)
cudnn.benchmark = True
# self.net = self.net.cuda()
# print('{} model, anchors, and classes loaded.'.format(self.Config["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 structure(self):
model = get_ssd("train", Config["num_classes"])
self.structure_start = summary()
self.structure_start.msg.connect(self.structure_msg)
self.structure_start.summary(model,
input_size=(3, Config["min_dim"],
Config["min_dim"]),
batch_size=Config["Batch_size"])
def generate(self):
self.num_classes = len(self.class_names) + 1
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ssd.get_ssd("test", self.num_classes)
self.net = model
model.load_state_dict(torch.load(self.model_path))
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))
# colors of boxes
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 generate(self):
# 计算总的种类
self.num_classes = len(self.class_names) + 1
# 载入模型，如果原来的模型里已经包括了模型结构则直接载入。
# 否则先构建模型再载入
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ssd.get_ssd("test", Config["num_classes"])
self.net = model
model.load_state_dict(torch.load(self.model_path))
if self.Cuda:
self.net = torch.nn.DataParallel(self.net)
cudnn.benchmark = True
self.net = self.net.cuda()
model.load_state_dict(torch.load(self.model_path))
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))
# ------------------------------------#
lr = 5e-4
freeze_lr = 1e-4
Cuda = True
Start_iter = 0
Freeze_epoch = 50
Epoch = 100
Batch_size = 4
#-------------------------------#
# Dataloder的使用
#-------------------------------#
Use_Data_Loader = True
model = get_ssd("train", Config["num_classes"])
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_data/ssd_weights.pth",
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!')
def normalize(I):
# 归一化梯度map，先归一化到 mean=0 std=1
norm = (I - I.mean()) / I.std()
# 把 std 重置为 0.1，让梯度map中的数值尽可能接近 0
norm = norm * 0.1
# 均值加 0.5，保证大部分的梯度值为正
norm = norm + 0.5
# 把 0，1 以外的梯度值分别设置为 0 和 1
norm = norm.clip(0, 1)
return norm
if __name__ == '__main__':
ssd = SSD()
model = get_ssd("train", 3) # ssd.net
model.load_state_dict(
torch.load(
"F:/Iris_SSD_small/ssd-pytorch-master/logs/Epoch50-Loc0.0260-Conf0.1510.pth",
map_location=torch.device('cuda')))
ssd.net = model.eval()
ssd.net = torch.nn.DataParallel(ssd.net)
# ssd.net = ssd.net.cpu() # ****
# criterion = MultiBoxLoss(3, 0.5, True, 0, True, 3, 0.5,False, True)
model = ssd.net.module
imgPath = '1.bmp'
image = Image.open(imgPath)
image.show()
image_size = image.size
image = image.convert('RGB')
def run(self):
Batch_size = Config["Batch_size"] # 每批次输入图片数量
lr = Config["lr"]
Epoch = Config["Epoch"]
Cuda = Config["Cuda"]
Start_iter = Config["Start_iter"]
model = get_ssd("train", Config["num_classes"])
self.msg.emit('将权重加载到字典中......')
model_dict = model.state_dict()
pretrained_dict = torch.load(Config['migrate_path'],
map_location=torch.device('cpu'))
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)
self.msg.emit('完成加载')
net = model
if Cuda:
net = torch.nn.DataParallel(model)
cudnn.benchmark = True
net = net.cuda()
annotation_path = 'neural_network/2007_train.txt'
with open(annotation_path) as f:
lines = f.readlines()
np.random.seed(10101)
np.random.shuffle(lines) # 打乱
np.random.seed(None)
num_train = len(lines)
# 生成图片和标签
gen = Generator(Batch_size, lines,
(Config["min_dim"], Config["min_dim"]),
Config["num_classes"]).generate()
# 优化器
optimizer = optim.Adam(net.parameters(), lr=lr)
criterion = MultiBoxLoss(Config['num_classes'], 0.5, True, 0, True, 3,
0.5, False, Cuda)
net.train()
epoch_size = num_train // Batch_size
for epoch in range(Start_iter, Epoch):
if epoch % 10 == 0:
self.adjust_learning_rate(optimizer, lr, 0.95, epoch)
self.loc_loss = 0
self.conf_loss = 0
for iteration in range(epoch_size):
if self.flag == True:
images, targets = next(gen)
with torch.no_grad():
if Cuda:
images = torch.from_numpy(images).cuda().type(
torch.FloatTensor)
targets = [
torch.from_numpy(ann).cuda().type(
torch.FloatTensor) 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
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
# 反向传播
loss.backward()
optimizer.step()
# 加上
self.loc_loss += loss_l.item()
self.conf_loss += loss_c.item()
# print('\nEpoch:' + str(epoch + 1) + '/' + str(Epoch))
self.msg.emit('Epoch:' + str(epoch + 1) + '/' + str(Epoch))
loss = '|| Loc_Loss: %.4f || Conf_Loss: %.4f ||' % (
self.loc_loss / (iteration + 1), self.conf_loss /
(iteration + 1))
self.msg.emit('iter:' + str(iteration) + '/' +
str(epoch_size) + str(loss))
else:
break
else:
if self.loc_loss / (iteration + 1) <= Config[
"loc_loss"] and self.conf_loss / (
iteration + 1) <= Config["conf_loss"]:
ep = 'neural_network/outputs/Epoch%d-Loc%.4f-Conf%.4f.pth' % (
epoch + 1), self.loc_loss / (
iteration + 1), self.conf_loss / (iteration + 1)
self.msg.emit('保存模型' + str(ep))
torch.save(
model.state_dict(),
'neural_network/outputs/Epoch%d-Loc%.4f-Conf%.4f.pth' %
((epoch + 1), self.loc_loss /
(iteration + 1), self.conf_loss / (iteration + 1)))
continue
self.msg.emit('训练中止')
break
else:
ep = 'neural_network/outputs/Epoch%d-Loc%.4f-Conf%.4f.pth' % (
(epoch + 1), self.loc_loss / (iteration + 1), self.conf_loss /
(iteration + 1))
self.msg.emit('保存最后模型' + str(ep))
torch.save(
model.state_dict(),
'neural_network/outputs/Epoch%d-Loc%.4f-Conf%.4f.pth' %
((epoch + 1), self.loc_loss / (iteration + 1), self.conf_loss /
(iteration + 1)))
self.finish.emit('训练完成')
