def test(config): is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) # net = net.cuda() net = net.to(device) # import torch.onnx # # 按照输入格式,设计随机输入 # dummy_input =torch.randn(1, 3, 416, 416).to(device) # # 导出模型 # torch.onnx.export(net.module,dummy_input, 'darknet.onnx',verbose=True) # exit() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("load checkpoint from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # prepare images path images_name = os.listdir(config["images_path"]) images_path = [ os.path.join(config["images_path"], name) for name in images_name ] if len(images_path) == 0: raise Exception("no image found in {}".format(config["images_path"])) # Start inference batch_size = config["batch_size"] for step in range(0, len(images_path), batch_size): # preprocess images = [] images_origin = [] for path in images_path[step * batch_size:(step + 1) * batch_size]: logging.info("processing: {}".format(path)) image = cv2.imread(path, cv2.IMREAD_COLOR) if image is None: logging.error("read path error: {}. skip it.".format(path)) continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) images_origin.append(image) # keep for save result image = cv2.resize(image, (config["img_w"], config["img_h"]), interpolation=cv2.INTER_LINEAR) image = image.astype(np.float32) image /= 255.0 image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) images.append(image) images = np.asarray(images) # images = torch.from_numpy(images).cuda() images = torch.from_numpy(images).to(device) # inference with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression( output, config["yolo"]["classes"], conf_thres=config["confidence_threshold"], nms_thres=0.45) # write result images. Draw bounding boxes and labels of detections classes = open(config["classes_names_path"], "r").read().split("\n")[:-1] if not os.path.isdir("./output/"): os.makedirs("./output/") for idx, detections in enumerate(batch_detections): plt.figure() fig, ax = plt.subplots(1) ax.imshow(images_origin[idx]) if detections is not None: unique_labels = detections[:, -1].cpu().unique() n_cls_preds = len(unique_labels) bbox_colors = random.sample(colors, n_cls_preds) for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: color = bbox_colors[int( np.where(unique_labels == int(cls_pred))[0])] # Rescale coordinates to original dimensions ori_h, ori_w = images_origin[idx].shape[:2] pre_h, pre_w = config["img_h"], config["img_w"] box_h = ((y2 - y1) / pre_h) * ori_h box_w = ((x2 - x1) / pre_w) * ori_w y1 = (y1 / pre_h) * ori_h x1 = (x1 / pre_w) * ori_w # Create a Rectangle patch bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=color, facecolor='none') # Add the bbox to the plot ax.add_patch(bbox) # Add label plt.text(x1, y1, s=classes[int(cls_pred)], color='white', verticalalignment='top', bbox={ 'color': color, 'pad': 0 }) # Save generated image with detections plt.axis('off') plt.gca().xaxis.set_major_locator(NullLocator()) plt.gca().yaxis.set_major_locator(NullLocator()) plt.savefig('output/{}_{}.jpg'.format(step, idx), bbox_inches='tight', pad_inches=0.0) plt.close() logging.info("Save all results to ./output/")
def train(config): config["global_step"] = config.get("start_step", 0) is_training = False if config.get("export_onnx") else True # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Optimizer and learning rate optimizer = _get_optimizer(config, net) lr_scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=config["lr"]["decay_step"], gamma=config["lr"]["decay_gamma"]) # Set data parallel net = nn.DataParallel(net) # net = net.cuda() net = net.to(device) # Restore pretrain model if config["pretrain_snapshot"]: logging.info("Load pretrained weights from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) # Only export onnx # if config.get("export_onnx"): # real_model = net.module # real_model.eval() # dummy_input = torch.randn(8, 3, config["img_h"], config["img_w"]).cuda() # save_path = os.path.join(config["sub_working_dir"], "pytorch.onnx") # logging.info("Exporting onnx to {}".format(save_path)) # torch.onnx.export(real_model, dummy_input, save_path, verbose=False) # logging.info("Done. Exiting now.") # sys.exit() # Evaluate interface # if config["evaluate_type"]: # logging.info("Using {} to evaluate model.".format(config["evaluate_type"])) # evaluate_func = importlib.import_module(config["evaluate_type"]).run_eval # config["online_net"] = net # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader dataloader = torch.utils.data.DataLoader(COCODataset( config["train_path"], (config["img_w"], config["img_h"]), is_training=True), batch_size=config["batch_size"], shuffle=True, num_workers=0, pin_memory=True) # Start the training loop logging.info("Start training.") for epoch in range(config["epochs"]): for step, samples in enumerate(dataloader): images, labels = samples["image"], samples["label"] start_time = time.time() config["global_step"] += 1 # Forward and backward optimizer.zero_grad() outputs = net(images) losses_name = ["total_loss", "x", "y", "w", "h", "conf", "cls"] losses = [] for _ in range(len(losses_name)): losses.append([]) for i in range(3): _loss_item = yolo_losses[i](outputs[i], labels) for j, l in enumerate(_loss_item): losses[j].append(l) losses = [sum(l) for l in losses] # print(losses) loss = losses[0] loss.backward() optimizer.step() if step > 0 and step % 10 == 0: _loss = loss.item() duration = float(time.time() - start_time) example_per_second = config["batch_size"] / duration lr = optimizer.param_groups[0]['lr'] logging.info( "epoch [%.3d] iter = %d loss = %.2f example/sec = %.3f lr = %.5f " % (epoch, step, _loss, example_per_second, lr)) config["tensorboard_writer"].add_scalar( "lr", lr, config["global_step"]) config["tensorboard_writer"].add_scalar( "example/sec", example_per_second, config["global_step"]) for i, name in enumerate(losses_name): value = _loss if i == 0 else losses[i] config["tensorboard_writer"].add_scalar( name, value, config["global_step"]) if step > 0 and step % 1000 == 0: # net.train(False) _save_checkpoint(net.state_dict(), config) # net.train(True) lr_scheduler.step() # net.train(False) _save_checkpoint(net.state_dict(), config) # net.train(True) logging.info("Bye~")