# 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
