def prepare_dataset_dataloader(args): if not os.path.exists(args.save_folder): os.mkdir(args.save_folder) dataset = COCODetection(image_path=cfg.dataset.train_images, info_file=cfg.dataset.train_info, transform=SSDAugmentation(MEANS)) val_dataset = None if args.validation_epoch > 0: setup_eval(args) val_dataset = COCODetection(image_path=cfg.dataset.valid_images, info_file=cfg.dataset.valid_info, transform=BaseTransform(MEANS)) # I don't use the timer during training (I use a different timing method). # Apparently there's a race condition with multiple GPUs, so disable it just to be safe. timer.disable_all() # Both of these can set args.resume to None, so do them before the check if args.resume == 'interrupt': args.resume = SavePath.get_interrupt(args.save_folder) elif args.resume == 'latest': args.resume = SavePath.get_latest(args.save_folder, cfg.name) data_loader = data.DataLoader(dataset, args.batch_size, num_workers=args.num_workers, shuffle=True, collate_fn=detection_collate, pin_memory=True) return dataset, val_dataset, data_loader
def compute_val_map(yolact_net): with torch.no_grad(): val_dataset = COCODetection(image_path=cfg.dataset.valid_images, info_file=cfg.dataset.valid_info, augmentation=BaseTransform()) yolact_net.eval() print("\nComputing validation mAP...", flush=True) table = evaluate(yolact_net, val_dataset, during_training=True) yolact_net.train() return table
def train(): if args.dataset == 'COCO': if args.dataset_root == VOC_ROOT: if not os.path.exists(COCO_ROOT): parser.error('Must specify dataset_root if specifying dataset') print("WARNING: Using default COCO dataset_root because " + "--dataset_root was not specified.") args.dataset_root = COCO_ROOT cfg = coco dataset = COCODetection(root=args.dataset_root, transform=SSDAugmentation(cfg['min_dim'], MEANS)) elif args.dataset == 'VOC': if args.dataset_root == COCO_ROOT: parser.error('Must specify dataset if specifying dataset_root') cfg = voc dataset = VOCDetection(root=args.dataset_root, transform=SSDAugmentation(cfg['min_dim'], MEANS)) elif args.dataset == 'CUSTOM': if args.dataset_root == VOC_ROOT or args.dataset_root == COCO_ROOT: parser.error('Must specify dataset if specifying dataset_root') cfg = custom dataset = customDetection(root=args.dataset_root, transform=SSDAugmentation(cfg['min_dim'], MEANS)) if args.visdom: import visdom viz = visdom.Visdom() ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes']) net = ssd_net if args.cuda: net = torch.nn.DataParallel(ssd_net) cudnn.benchmark = True if args.resume: print('Resuming training, loading {}...'.format(args.resume)) ssd_net.load_weights(args.resume) else: pass # resnet_weights = torch.load(args.save_folder + args.basenet) # print('Loading base network...') # ssd_net.resnet.load_state_dict(resnet_weights) if args.cuda: net = net.cuda() if not args.resume: print('Initializing weights...') # initialize newly added layers' weights with xavier method ssd_net.extras.apply(weights_init) ssd_net.loc.apply(weights_init) ssd_net.conf.apply(weights_init) optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5, False, args.cuda) net.train() # loss counters loc_loss = 0 conf_loss = 0 epoch = 0 print('Loading the dataset...') epoch_size = len(dataset) // args.batch_size print('Epochj Size:', epoch_size) print('Training SSD on:', dataset.name) print('Using the specified args:') print(args) step_index = 0 if args.visdom: vis_title = 'SSD.PyTorch on ' + dataset.name vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss'] iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend) epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend) data_loader = data.DataLoader(dataset, args.batch_size, num_workers=args.num_workers, shuffle=True, collate_fn=detection_collate, pin_memory=True) # create batch iterator batch_iterator = iter(data_loader) for iteration in range(args.start_iter, cfg['max_iter']): if args.visdom and iteration != 0 and (iteration % epoch_size == 0): update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot, None, 'append', epoch_size) # reset epoch loss counters loc_loss = 0 conf_loss = 0 epoch += 1 if iteration in cfg['lr_steps']: step_index += 1 adjust_learning_rate(optimizer, args.gamma, step_index) # load train data # images, targets = next(batch_iterator) try: images, targets = next(batch_iterator) except StopIteration: batch_iterator = iter(data_loader) images, targets = next(batch_iterator) if args.cuda: images = Variable(images.cuda()) targets = [Variable(ann.cuda(), volatile=True) for ann in targets] else: images = Variable(images) targets = [Variable(ann, volatile=True) for ann in targets] # forward t0 = time.time() out = net(images) # backprop optimizer.zero_grad() loss_l, loss_c = criterion(out, targets) loss = loss_l + loss_c loss.backward() optimizer.step() t1 = time.time() # loc_loss += loss_l.data[0] # conf_loss += loss_c.data[0] loc_loss += loss_l.item() conf_loss += loss_c.item() if iteration % 10 == 0: print('timer: %.4f sec.' % (t1 - t0)) # print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ') print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.item()), end=' ') if args.visdom: # update_vis_plot(iteration, loss_l.data[0], loss_c.data[0], # iter_plot, epoch_plot, 'append') update_vis_plot(iteration, loss_l.item(), loss_c.item(), iter_plot, epoch_plot, 'append') if iteration != 0 and iteration % 5000 == 0: print('Saving state, iter:', iteration) torch.save(ssd_net.state_dict(), args.save_folder + '/ssd300_COCO_' + repr(iteration) + '.pth') torch.save(ssd_net.state_dict(), args.save_folder + '' + args.dataset + '.pth')
strs = args.trained_model.split('_') config = f'{strs[-3]}_{strs[-2]}_config' update_config(config) print(f'\nUsing \'{config}\' according to the trained_model.\n') with torch.no_grad(): if cuda: cudnn.benchmark = True cudnn.fastest = True torch.set_default_tensor_type('torch.cuda.FloatTensor') else: torch.set_default_tensor_type('torch.FloatTensor') dataset = COCODetection(cfg.dataset.valid_images, cfg.dataset.valid_info, augmentation=BaseTransform()) net = Yolact() # net=onnx.load("yolact.onnx") # net= torch.jit.script(net) net.load_weights('weights/' + args.trained_model, cuda) net.eval() if (args.onnx): ONNX_util.save_yolact(net, dataset, "yolact.onnx") print('\nModel loaded.\n') if cuda: net = net.cuda() evaluate(net, dataset, args.max_num, False, args.cocoapi,
if args.cuda: cudnn.benchmark = True cudnn.fastest = True torch.set_default_tensor_type('torch.cuda.FloatTensor') else: torch.set_default_tensor_type('torch.FloatTensor') if args.resume and not args.display: with open(args.ap_data_file, 'rb') as f: ap_data = pickle.load(f) calc_map(ap_data) exit() if args.image is None and args.video is None and args.images is None: dataset = COCODetection(cfg.dataset.valid_images, cfg.dataset.valid_info, transform=BaseTransform(), has_gt=cfg.dataset.has_gt) prep_coco_cats() else: dataset = None print('Loading model...', end='') net = Yolact() net.load_weights(args.trained_model) net.eval() print(' Done.') if args.cuda: net = net.cuda() net.detect.use_fast_nms = args.fast_nms
optimizer = optim.SGD(net.parameters(), lr=cfg.lr, momentum=cfg.momentum, weight_decay=cfg.decay) criterion = Multi_Loss(num_classes=cfg.num_classes, pos_thre=cfg.pos_iou_thre, neg_thre=cfg.neg_iou_thre, np_ratio=3) if cuda: cudnn.benchmark = True net = nn.DataParallel(net).cuda() criterion = nn.DataParallel(criterion).cuda() dataset = COCODetection(image_path=cfg.dataset.train_images, info_file=cfg.dataset.train_info, augmentation=SSDAugmentation()) data_loader = data.DataLoader(dataset, cfg.batch_size, num_workers=8, shuffle=True, collate_fn=detection_collate, pin_memory=True) step_index = 0 start_step = resume_step if args.resume else 0 batch_time = MovingAverage() loss_types = ['B', 'C', 'M', 'S'] loss_avgs = {k: MovingAverage() for k in loss_types} map_tables = []
# elif args.benchmark: # print() # print() # print('Stats for the last frame:') # timer.print_stats() # avg_seconds = frame_times.get_avg() # print('Average: %5.2f fps, %5.2f ms' % (1 / frame_times.get_avg(), 1000 * avg_seconds)) if __name__ == '__main__': # 数据集与标签 valid_dataset = COCODetection(image_path='./data/coco/images/val2017/', info_file='./data/coco/annotations/instances_val2017.json', transform=BaseTransform(), has_gt=True ) prep_coco_cats() # 模型 print('Loading model...', end='') model = Yolact() model.load_weights(args.trained_model) model.eval() model = model.cuda() if args.cuda else model.cpu() print(' Done.') # 核心入口 with torch.no_grad(): if not os.path.exists('results'):
if __name__ == '__main__': args = arg_parse() weightsfile = args.weights nms_thresh = args.nms_thresh input_wh = args.input_wh cuda = args.cuda use_pad = args.use_pad save_folder = args.save_folder dataset = args.dataset if dataset[0] == "V": cfg = voc_config test_dataset = VOCDetection(cfg["root"], datasets_dict["VOC2007"], input_wh) elif dataset[0] == "C": cfg = coco_config test_dataset = COCODetection(cfg["root"], datasets_dict["COCOval"], input_wh) else: print("only support VOC and COCO datasets !!!") print("load test_dataset successfully.....") with open(cfg["name_path"], "r") as f: classes = [i.strip() for i in f.readlines()] net = Yolov3("test", input_wh, cfg["anchors"], cfg["anchors_mask"], cfg["num_classes"]) state_dict = torch.load(weightsfile) from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): head = k[:7]
help='visible devices for CUDA') args = parser.parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus if torch.cuda.is_available(): torch.set_default_tensor_type('torch.cuda.FloatTensor') dataset = None if args.dataset == 'VOC': config = voc rt = VOC_ROOT if args.dataset_root is None else args.dataset_root dataset = VOCDetection(rt, transform=BaseTransform(300, (104, 117, 123))) elif args.dataset == 'COCO18': config = coco18 rt = COCO_ROOT if args.dataset_root is None else args.dataset_root dataset = COCODetection(rt, transform=BaseTransform(300, (104, 117, 123)), target_transform=COCOAnnotationTransform('COCO18')) elif args.dataset == 'helmet': config = helmet rt = HELMET_ROOT if args.dataset_root is None else args.dataset_root dataset = HelmetDetection(rt, transform=BaseTransform(300, (104, 117, 123))) else: raise NotImplementedError() config = generic if __name__ == '__main__': params = torch.load(args.priors_pth) prior_types = trim(params) pool = AnchorsPool(dataset, prior_types,
def train(): if args.dataset == 'COCO18': # cfg = coco18 cfg = vococo rt = args.dataset_root or COCO_ROOT dataset = COCODetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS), target_transform=COCOAnnotationTransform('COCO18')) elif args.dataset == 'COCO': cfg = coco # cfg = vococo rt = args.dataset_root or COCO_ROOT dataset = COCODetection(root=rt, image_sets=(('2017', 'train'),), transform=SSDAugmentation(cfg['min_dim'], MEANS)) elif args.dataset == 'VOC': cfg = voc # cfg = coco_on_voc rt = args.dataset_root or VOC_ROOT dataset = VOCDetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS)) elif args.dataset == 'BCCD': cfg = bccd rt = args.dataset_root or BCCD_ROOT dataset = BCCDDetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS)) elif args.dataset == 'SHWD': cfg = shwd rt = args.dataset_root or SHWD_ROOT dataset = SHWDDetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS)) elif args.dataset == 'helmet': cfg = helmet rt = args.dataset_root or HELMET_ROOT dataset = HelmetDetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS)) else: raise RuntimeError() if args.custom_priors is not None: apt = IOAdapterSSD(cfg, 'test') apt.load(*torch.load(args.custom_priors)) custom_priors = apt.fit_output(apt.msks[0]) print('num_boxes = %d ' % custom_priors.size()[0]) custom_mbox = None # params = torch.load(args.custom_priors) # # bbox = gen_priors(params, args.prior_types, cfg) # gen = AdaptivePriorBox(cfg, phase='test') # custom_priors = gen.forward(params) # custom_mbox = [p.size(0) for p in params] if args.cuda: custom_priors = custom_priors.cuda() ssd_net = build_ssd('train', cfg, custom_mbox, custom_priors) else: # priors = torch.load('anchors/voc_baseline.pth') # if args.cuda: # priors = priors.cuda() ssd_net = build_ssd('train', cfg) net = ssd_net if args.cuda: net = torch.nn.DataParallel(ssd_net).cuda() cudnn.benchmark = True if args.resume: print('Resuming training, loading {}...'.format(args.resume)) ssd_net.load_weights(args.resume) else: vgg_weights = torch.load(args.save_folder + args.basenet) print('Loading base network...') ssd_net.vgg.load_state_dict(vgg_weights) if not args.resume: print('Initializing weights...') # initialize newly added layers' weights with xavier method ssd_net.extras.apply(weights_init) ssd_net.loc.apply(weights_init) ssd_net.conf.apply(weights_init) optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5, False, best_prior_weight=args.k, use_gpu=args.cuda) net.train() # loss counters loc_loss = 0 conf_loss = 0 epoch = 0 print('Loading the dataset...') epoch_size = len(dataset) // args.batch_size print('Training SSD on:', dataset.name) print('Using the specified args:') print(args) step_index = 0 data_loader = data.DataLoader(dataset, args.batch_size, num_workers=args.num_workers, shuffle=True, collate_fn=detection_collate, pin_memory=True) # create batch iterator batch_iterator = iter(data_loader) for step in cfg['lr_steps']: if args.start_iter > step: print('over %d steps, adjust lr' % step) step_index += 1 adjust_learning_rate(optimizer, args.gamma, step_index) else: break for iteration in range(args.start_iter, cfg['max_iter']): if iteration in cfg['lr_steps']: step_index += 1 adjust_learning_rate(optimizer, args.gamma, step_index) # load train data try: images, targets = next(batch_iterator) except StopIteration: batch_iterator = iter(data_loader) images, targets = next(batch_iterator) if args.cuda: images = images.cuda() targets = [ann.cuda() for ann in targets] # targets = targets.cuda() # else: # # targets = [ann for ann in targets] # forward t0 = time.time() out = net(images) # backprop optimizer.zero_grad() loss_l, loss_c = criterion(out, targets) loss = loss_l + loss_c loss.backward() optimizer.step() t1 = time.time() loc_loss += loss_l.item() conf_loss += loss_c.item() if iteration % 10 == 0: print('timer: %.4f sec.' % (t1 - t0)) print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.item()), end=' ') if iteration != 0 and iteration % 2000 == 0: print('Saving state, iter:', iteration) torch.save(ssd_net.state_dict(), ('weights/cache/%s_%s_' % (args.save_name, args.dataset)) + repr(iteration) + '.pth') name = '%s_%s' % (args.save_name, args.dataset) torch.save(ssd_net.state_dict(), args.save_folder + name + '.pth')
"COCO": [('2014', 'train'), ('2014', 'valminusminival')], "VOC2007": [('0712', "2007_test")], "COCOval": [('2014', 'minival')] } print('Loading Dataset...') if dataset_name[0] == "V": cfg = voc_config train_dataset = VOCDetection(cfg["root"], datasets_version[dataset_name], input_wh, batch_size, cfg["multiscale"], dataset_name) elif dataset_name[0] == "C": cfg = coco_config train_dataset = COCODetection(cfg["root"], datasets_version[dataset_name], input_wh, batch_size, cfg["multiscale"], dataset_name) else: print('Unkown dataset!') # load Yolov3 net net = Yolov3("train", input_wh, cfg["anchors"], cfg["anchors_mask"], cfg["num_classes"]) if args.resume_net == None: net.load_weights(basenet) else: state_dict = torch.load(args.resume_net) from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): head = k[:7]
self.ptr = 0 self.size = size self.ids = [i for i in range(lst.size(0))] if shuffle: random.shuffle(self.ids) if drop_last: self.ids = self.ids[:int(lst.size(0) / size) * size] def __next__(self): self.ptr += self.size if self.ptr > len(self.lst): raise StopIteration() return self.lst[self.ids[self.ptr - self.size:self.ptr]].clone().detach() def __iter__(self): return self.IterInstance(self.bb_data, self.b_size, self.shuffle, self.drop_last) if __name__ == '__main__': rt = COCO_ROOT data_set = COCODetection( root=rt, image_sets=(('2017', 'train'), ), transform=BaseTransform(300, (104, 117, 123)), target_transform=COCOAnnotationTransform(keep_difficult=False)) loader = BoundingBoxesLoader(data_set, [i for i in range(len(COCO_CLASSES))], cache_pth='../truths/gts_coco_17train.pth')
VOCAnnotationTransform(), no_anno=no_anno) elif args.dataset == 'SHWD': dataset = SHWDDetection(root, (set_type, ), BaseTransform(300, dataset_mean)) # elif args.dataset == 'VOC07': # dataset = VOCDetection(root, [('2007', set_type)], # BaseTransform(300, dataset_mean), # VOCAnnotationTransform()) # elif args.dataset == 'VOC-v2': # dataset = VOCDetection(root, [('2012', set_type)], # BaseTransform(300, dataset_mean), # VOCAnnotationTransform()) elif args.dataset == 'COCO18': dataset = COCODetection(root, [('18', set_type)], BaseTransform(300, dataset_mean), COCOAnnotationTransform('COCO18')) elif args.dataset == 'COCO': dataset = COCODetection(root, [(year, set_type)], BaseTransform(300, dataset_mean), no_anno=no_anno) if args.cuda: net = net.cuda() cudnn.benchmark = True # evaluation test_net(args.save_folder, net, args.cuda, dataset, BaseTransform(net.size, dataset_mean),
file, use_gpu=use_gpu, use_fast_prep=False) ap_data = { 'box': [[APDataObject() for _ in COCO_CLASSES] for _ in iou_thresholds], 'mask': [[APDataObject() for _ in COCO_CLASSES] for _ in iou_thresholds] } valid_images = '../COCO/val2017/' valid_info = '../COCO/annotations/instances_val2017.json' transform = get_transform(backbone_name) val_dataset = COCODetection(image_path=valid_images, info_file=valid_info, transform=BaseTransform( transform, 550, MEANS, STD)) # 不使用数据增强 num_imgs = len(val_dataset) start = time.time() # 不管训练集还是验证集,坐标gt(x0, y0, x1, y1)都是归一化后的值。 # 训练集的gt_masks resize成550x550,而验证集的gt_masks没有做resize。 for image_idx in range(num_imgs): img, gt, gt_masks, num_crowd = val_dataset.pull_item(image_idx) _, h, w = gt_masks.shape # 因为验证集的gt_masks没有做resize。所以从这里获得原图高宽 boxes, masks, classes, scores = _decode.eval_image(img, h, w) if len(boxes) == 0: continue box_scores = scores mask_scores = scores masks = np.reshape(masks, (-1, h * w))
def train(config, model): if config.name == 'COCO': from data.coco import COCODetection, COCOAnnotationTransform, COCO_CLASSES, get_label_map elif config.name == 'VOC': from data.voc0712 import VOCDetection, VOCAnnotationTransform, VOC_CLASSES if not os.path.exists(config.train['save_folder']): os.mkdir(config.train['save_folder']) if config.train['dataset'] == 'COCO': dataset = COCODetection(root=config.train['dataset_root'], transform=SSDAugmentation( config.min_dim, config.MEANS)) elif config.train['dataset'] == 'VOC': dataset = VOCDetection(root=config.train['dataset_root'], transform=SSDAugmentation( config.min_dim, config.MEANS)) if config.train['visdom']: import visdom viz = visdom.Visdom() if config.train['cuda']: net = torch.nn.DataParallel(model.cuda()) cudnn.benchmark = True if config.train['resume']: print('Resuming training, loading {}...'.format( config.train['resume'])) # model.load_weights(config.train['resume']) optimizer = optim.SGD(net.parameters(), lr=config.train['lr'], momentum=config.train['momentum'], weight_decay=config.train['weight_decay']) criterion = MultiBoxLoss(config.variance, config.num_classes, 0.5, True, 0, True, 3, 0.5, False, config.train['cuda']) lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer=optimizer, milestones=config.train['milestones']) epochs = config.train['milestones'][-1] net.train() loc_loss = 0 conf_loss = 0 epoch = 0 print('Loading the dataset...') epoch_size = len(dataset) // config.train['batch_size'] print('Training SSD on:', dataset.name) step_index = 0 if config.train['visdom']: vis_title = 'SSD.PyTorch on ' + dataset.name vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss'] iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend) epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend) data_loader = data.DataLoader(dataset, config.train['batch_size'], num_workers=config.train['num_workers'], shuffle=True, collate_fn=detection_collate, pin_memory=True) print(len(dataset)) global_step = 0 for epoch in range(epochs): print(config.train['save_folder'] + '' + config.train['dataset'] + str(epoch) + '.pth') for step, (images, targets) in enumerate(data_loader): print(targets) exit(0) if config.train['cuda']: images = images.cuda() targets = [ann.cuda() for ann in targets] else: images = images targets = [ann for ann in targets] # forward t0 = time.time() out = net(images) # backprop loss_l, loss_c = criterion(out, targets) loss = loss_l + loss_c optimizer.zero_grad() loss.backward() optimizer.step() t1 = time.time() loc_loss += loss_l.item() conf_loss += loss_c.item() if step > 0 and step % 10 == 0: print('timer: %.4f sec.' % (t1 - t0)) print('epoch' + repr(epoch) + ' iter ' + repr(step) + ' || Loss: %.4f ||' % (loss.item()), end=' ') # if config.train['visdom']: # update_vis_plot(iteration, loss_l.data[0], loss_c.data[0], iter_plot, epoch_plot, 'append') if step != 0 and step % 5000 == 0: print('Saving state, iter:', global_step) torch.save( model.state_dict(), 'weights/peleenet_ssd800_COCO_' + repr(global_step) + '.pth') global_step += 1 # adjust lr lr_scheduler.step() torch.save( model.state_dict(), config.train['save_folder'] + '' + config.train['dataset'] + str(epoch) + '.pth')
# 从头训练。若使用focal_loss,预测的概率初始化为0.01 if use_focal_loss: prior = 0.01 b = -math.log((1.0 - prior) / prior) c_w = fluid.global_scope().find_var( 'prediction_layers.0.conf_layer.weight').get_tensor() c_b = fluid.global_scope().find_var( 'prediction_layers.0.conf_layer.bias').get_tensor() src_w = np.array(c_w) c_w.set((np.zeros(src_w.shape)).astype(np.float32), place) src_b = np.array(c_b) c_b.set((np.ones(src_b.shape) * b).astype(np.float32), place) # 验证集和训练集 dataset = COCODetection(image_path=train_images_path, info_file=train_anno_path, transform=SSDAugmentation(transform, input_size, MEANS, STD)) if eval: val_dataset = COCODetection(image_path=valid_images_path, info_file=valid_anno_path, transform=BaseTransform( transform, input_size, MEANS, STD)) # 不使用数据增强 num_train = len(dataset) if eval: num_val = len(val_dataset) # 一轮的步数 train_steps = int(num_train / batch_size) if eval: val_steps = int(num_val / batch_size)