# Resume if opt.resume: # resume an interrupted run ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path log_dir = Path(ckpt).parent.parent # runs/exp0 assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist' with open(log_dir / 'opt.yaml') as f: opt = argparse.Namespace(**yaml.load(f, Loader=yaml.FullLoader)) # replace opt.cfg, opt.weights, opt.resume = '', ckpt, True logger.info('Resuming training from %s' % ckpt) else: # opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml') opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(opt.cfg), check_file(opt.hyp) # check files assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified' opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size))) # extend to 2 sizes (train, test) log_dir = increment_dir(Path(opt.logdir) / 'exp', opt.name) # runs/exp1 # DDP mode device = select_device(opt.device, batch_size=opt.batch_size) if opt.local_rank != -1: assert torch.cuda.device_count() > opt.local_rank torch.cuda.set_device(opt.local_rank) device = torch.device('cuda', opt.local_rank) dist.init_process_group(backend='nccl', init_method='env://') # distributed backend assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count' opt.batch_size = opt.total_batch_size // opt.world_size # Hyperparameters with open(opt.hyp) as f: hyp = yaml.load(f, Loader=yaml.FullLoader) # load hyps if 'box' not in hyp:
def detect(save_img=False): save_dir, source, weights, view_img, save_txt, imgsz = \ Path(opt.save_dir), opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size webcam = source.isnumeric() or source.endswith('.txt') or \ source.lower().startswith(('rtsp://', 'rtmp://', 'http://')) # Directories if save_dir == Path('runs/detect'): # if default save_dir.mkdir(parents=True, exist_ok=True) # make base save_dir = Path(increment_dir(save_dir / 'exp', opt.name)) # increment run (save_dir / 'labels' if save_txt else save_dir).mkdir( parents=True, exist_ok=True) # make new dir # Initialize set_logging() device = select_device(opt.device) half = device.type != 'cpu' # half precision only supported on CUDA # Load model model = attempt_load(weights, map_location=device) # load FP32 model imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size if half: model.half() # to FP16 # Second-stage classifier classify = False if classify: modelc = load_classifier(name='resnet101', n=2) # initialize modelc.load_state_dict( torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights modelc.to(device).eval() # Set Dataloader vid_path, vid_writer = None, None if webcam: view_img = True cudnn.benchmark = True # set True to speed up constant image size inference dataset = LoadStreams(source, img_size=imgsz) else: save_img = True dataset = LoadImages(source, img_size=imgsz) # Get names and colors names = model.module.names if hasattr(model, 'module') else model.names colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))] # Run inference t0 = time.time() img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img _ = model(img.half() if half else img ) if device.type != 'cpu' else None # run once for path, img, im0s, vid_cap in dataset: img = torch.from_numpy(img).to(device) img = img.half() if half else img.float() # uint8 to fp16/32 img /= 255.0 # 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0) # Inference t1 = time_synchronized() pred = model(img, augment=opt.augment)[0] # Apply NMS pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms) t2 = time_synchronized() # Apply Classifier if classify: pred = apply_classifier(pred, modelc, img, im0s) # Process detections for i, det in enumerate(pred): # detections per image if webcam: # batch_size >= 1 p, s, im0 = Path(path[i]), '%g: ' % i, im0s[i].copy() else: p, s, im0 = Path(path), '', im0s save_path = str(save_dir / p.name) txt_path = str(save_dir / 'labels' / p.stem) + ( '_%g' % dataset.frame if dataset.mode == 'video' else '') s += '%gx%g ' % img.shape[2:] # print string gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh if det is not None and len(det): # Rescale boxes from img_size to im0 size det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() # Print results for c in det[:, -1].unique(): n = (det[:, -1] == c).sum() # detections per class s += '%g %ss, ' % (n, names[int(c)]) # add to string # Write results numCount = 0 for *xyxy, conf, cls in reversed(det): numCount += 1 if save_txt: # Write to file xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh line = (cls, *xywh, conf) if opt.save_conf else ( cls, *xywh) # label format with open(txt_path + '.txt', 'a') as f: f.write(('%g ' * len(line) + '\n') % line) if save_img or view_img: # Add bbox to image label = '%s %.2f' % (names[int(cls)], conf) plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) # Print time (inference + NMS) print('%sDone. (%.3fs)' % (s, t2 - t1)) print("Num of plants in the image:", numCount) # Stream results if view_img: cv2.imshow(p, im0) if cv2.waitKey(1) == ord('q'): # q to quit raise StopIteration # Save results (image with detections) if save_img: if dataset.mode == 'images': cv2.imwrite(save_path, im0) else: if vid_path != save_path: # new video vid_path = save_path if isinstance(vid_writer, cv2.VideoWriter): vid_writer.release( ) # release previous video writer fourcc = 'mp4v' # output video codec fps = vid_cap.get(cv2.CAP_PROP_FPS) w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) vid_writer = cv2.VideoWriter( save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) vid_writer.write(im0) if save_txt or save_img: print('Results saved to %s' % save_dir) print('Done. (%.3fs)' % (time.time() - t0))
def __init__(self, load_ext_mem=False): self.logger = logging.getLogger(__name__) self.device = torch.device('cuda:0') opt, hyp = get_opt_and_hyp() self.logger.info(f'Hyperparameters {hyp}') # self.log_dir = Path(opt.logdir) # logging directory self.log_dir = increment_dir(Path(opt.logdir) / 'exp', opt.name) # runs/exp1 self.log_dir = Path(self.log_dir) wdir = self.log_dir / 'weights' # weights dfirectory os.makedirs(wdir, exist_ok=True) # last = wdir / 'last.pt' # best = wdir / 'best.pt' # results_file = str(log_dir / 'results.txt') # epochs, batch_size, total_batch_size, weights, rank = \ # opt.epochs_init, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank # Save run settings with open(self.log_dir / 'hyp.yaml', 'w') as f: yaml.dump(hyp, f, sort_keys=False) with open(self.log_dir / 'opt.yaml', 'w') as f: yaml.dump(vars(opt), f, sort_keys=False) # Configure init_seeds(2) # with open(opt.data) as f: # data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict # check_dataset(data_dict) # check self.nc = 6 self.names = ['???'] * 6 # assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check # Model weights = opt.weights pretrained = weights.endswith('.pt') if pretrained: attempt_download(weights) # download if not found locally ckpt = torch.load(weights, map_location=self.device) # load checkpoint if hyp.get('anchors'): ckpt['model'].yaml['anchors'] = round(hyp['anchors']) # force autoanchor self.model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=self.nc).to(self.device) # create exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [] # exclude keys state_dict = ckpt['model'].float().state_dict() # to FP32 state_dict = intersect_dicts(state_dict, self.model.state_dict(), exclude=exclude) # intersect self.model.load_state_dict(state_dict, strict=False) # load self.logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(self.model.state_dict()), weights)) # report else: self.model = Model(opt.cfg, ch=3, nc=self.nc).to(self.device) # create # Freeze # freeze = [] # parameter names to freeze (full or partial) # for k, v in self.model.named_parameters(): # v.requires_grad = True # train all layers # if any(x in k for x in freeze): # print('freezing %s' % k) # v.requires_grad = False # Optimizer nbs = 64 # nominal batch size accumulate = max(round(nbs / opt.total_batch_size), 1) # accumulate loss before optimizing hyp['weight_decay'] *= opt.total_batch_size * accumulate / nbs # scale weight_decay pg0, pg1, pg2 = [], [], [] # optimizer parameter groups for k, v in self.model.named_modules(): if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): pg2.append(v.bias) # biases if isinstance(v, nn.BatchNorm2d): pg0.append(v.weight) # no decay elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): pg1.append(v.weight) # apply decay self.optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) self.optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay self.optimizer.add_param_group({'params': pg2}) # add pg2 (biases) #logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0))) del pg0, pg1, pg2 # Scheduler https://arxiv.org/pdf/1812.01187.pdf # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR # lf = lambda x: ((1 + math.cos(x * math.pi / opt.epochs)) / 2) * (1 - hyp['lrf']) + hyp['lrf'] # cosine #self.scheduler = lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs) # Resume # start_epoch = 0 self.best_fitness = 0.0 if pretrained: # Optimizer if ckpt['optimizer'] is not None: self.optimizer.load_state_dict(ckpt['optimizer']) self.best_fitness = ckpt['best_fitness'] # Epochs # start_epoch = ckpt['epoch'] + 1 del ckpt, state_dict # Image sizes self.gs = int(max(self.model.stride)) # grid size (max stride) # Model parameters hyp['cls'] *= self.nc / 80. # scale coco-tuned hyp['cls'] to current dataset self.model.nc = self.nc # attach number of classes to model self.model.hyp = hyp # attach hyperparameters to model self.model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou) # extenral memory self.extMem = externalMemory(size=200) print(f'external memory file: {self.extMem.get_memory_file()}') self.opt = opt self.hyp = hyp self.added_classes = 0
if args.cache_images: args.workers = 1 else: args.workers = os.cpu_count() args.workers = min(os.cpu_count(), args.workers) if args.test_batch_size < 0: args.test_batch_size = args.batch_size # 打印参数 logger.info("args: %s" % args) # 预定义的参数(不打印) args.img_size = (94, 24) args.lpr_max_len = 9 * 2 # 车牌最大位数 * 2 args.lpr_dropout_rate = .5 # 自动调整的参数(不打印) args.cache_dir = os.path.join(args.worker_dir, 'cache') args.out_dir = increment_dir(Path(args.worker_dir) / 'exp') args.weights_dir = os.path.join(args.out_dir, 'weights') # 参数处理后的初始化工作 os.makedirs(args.cache_dir, exist_ok=True) os.makedirs(args.weights_dir, exist_ok=True) set_cache_dir(args.cache_dir) load_ccpd_imgs(args.source_dir) # 开始训练 main(args)
torch.cuda.set_device(opt.local_rank) device = torch.device('cuda', opt.local_rank) dist.init_process_group(backend='nccl', init_method='env://') # distributed backend opt.world_size = dist.get_world_size() assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count' opt.batch_size = opt.total_batch_size // opt.world_size print(opt) # Train if not opt.evolve: tb_writer = None if opt.local_rank in [-1, 0]: print('Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/') tb_writer = SummaryWriter(log_dir=increment_dir('runs/exp', opt.name)) train(hyp, opt, device, tb_writer) # Evolve hyperparameters (optional) else: # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit) meta = {'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3) 'momentum': (0.1, 0.6, 0.98), # SGD momentum/Adam beta1 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay 'giou': (1, 0.02, 0.2), # GIoU loss gain 'cls': (1, 0.2, 4.0), # cls loss gain 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels) 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight 'iou_t': (0, 0.1, 0.7), # IoU training threshold