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('训练完成')