def _train(dataset_name: str, backbone_name: str, path_to_data_dir: str, path_to_checkpoints_dir: str, path_to_resuming_checkpoint: Optional[str]): dataset = DatasetBase.from_name(dataset_name)(path_to_data_dir, DatasetBase.Mode.TRAIN, Config.IMAGE_MIN_SIDE, Config.IMAGE_MAX_SIDE) dataloader = DataLoader(dataset, batch_size=Config.BATCH_SIZE, sampler=DatasetBase.NearestRatioRandomSampler( dataset.image_ratios, num_neighbors=Config.BATCH_SIZE), num_workers=0, collate_fn=DatasetBase.padding_collate_fn, pin_memory=True) #为便于调试,num_works置为0 Log.i('Found {:d} samples'.format(len(dataset))) backbone = BackboneBase.from_name(backbone_name)(pretrained=True) model = nn.DataParallel( Model(backbone, dataset.num_classes(), pooler_mode=Config.POOLER_MODE, anchor_ratios=Config.ANCHOR_RATIOS, anchor_sizes=Config.ANCHOR_SIZES, rpn_pre_nms_top_n=Config.RPN_PRE_NMS_TOP_N, rpn_post_nms_top_n=Config.RPN_POST_NMS_TOP_N, anchor_smooth_l1_loss_beta=Config.ANCHOR_SMOOTH_L1_LOSS_BETA, proposal_smooth_l1_loss_beta=Config.PROPOSAL_SMOOTH_L1_LOSS_BETA ).cuda()) # 便于调试 # model = Model( # backbone, dataset.num_classes(), pooler_mode=Config.POOLER_MODE, # anchor_ratios=Config.ANCHOR_RATIOS, anchor_sizes=Config.ANCHOR_SIZES, # rpn_pre_nms_top_n=Config.RPN_PRE_NMS_TOP_N, rpn_post_nms_top_n=Config.RPN_POST_NMS_TOP_N, # anchor_smooth_l1_loss_beta=Config.ANCHOR_SMOOTH_L1_LOSS_BETA, proposal_smooth_l1_loss_beta=Config.PROPOSAL_SMOOTH_L1_LOSS_BETA # ).cuda() ''' 训练用参数: IMAGE_MIN_SIDE: float = 600.0 IMAGE_MAX_SIDE: float = 1000.0 ANCHOR_RATIOS: List[Tuple[int, int]] = [(1, 2), (1, 1), (2, 1)] ANCHOR_SIZES: List[int] = [128, 256, 512] POOLER_MODE: Pooler.Mode = Pooler.Mode.ALIGN RPN_PRE_NMS_TOP_N: int = 12000 RPN_POST_NMS_TOP_N: int = 2000 ANCHOR_SMOOTH_L1_LOSS_BETA: float = 1.0 PROPOSAL_SMOOTH_L1_LOSS_BETA: float = 1.0 BATCH_SIZE: int = 1 LEARNING_RATE: float = 0.001 MOMENTUM: float = 0.9 WEIGHT_DECAY: float = 0.0005 STEP_LR_SIZES: List[int] = [50000, 70000] STEP_LR_GAMMA: float = 0.1 WARM_UP_FACTOR: float = 0.3333 WARM_UP_NUM_ITERS: int = 500 NUM_STEPS_TO_DISPLAY: int = 20 NUM_STEPS_TO_SNAPSHOT: int = 10000 NUM_STEPS_TO_FINISH: int = 90000 ''' #动量的意义: #1.降低病态条件数带来的振荡 #2.减少随机梯度带来的方差(权值的衰减也有这个用处) #优化算法的两种衰减: #1. 权值的衰减:表现为在总的损失函数后面再加上权值的L2范数 #2.学习率的衰减:表现为通过学习率调节器以不同策略随着学习步增加,对学习率进行衰减调节 #optimizer = optim.Adam(model.parameters()) optimizer = optim.SGD(model.parameters(), lr=Config.LEARNING_RATE, momentum=Config.MOMENTUM, weight_decay=Config.WEIGHT_DECAY) scheduler = WarmUpMultiStepLR(optimizer, milestones=Config.STEP_LR_SIZES, gamma=Config.STEP_LR_GAMMA, factor=Config.WARM_UP_FACTOR, num_iters=Config.WARM_UP_NUM_ITERS) step = 0 time_checkpoint = time.time() losses = deque(maxlen=100) summary_writer = SummaryWriter( os.path.join(path_to_checkpoints_dir, 'summaries')) should_stop = False num_steps_to_display = Config.NUM_STEPS_TO_DISPLAY num_steps_to_snapshot = Config.NUM_STEPS_TO_SNAPSHOT num_steps_to_finish = Config.NUM_STEPS_TO_FINISH if path_to_resuming_checkpoint is not None: step = model.module.load(path_to_resuming_checkpoint, optimizer, scheduler) Log.i( f'Model has been restored from file: {path_to_resuming_checkpoint}' ) device_count = torch.cuda.device_count() #BATCH_SIZE默认是1 assert Config.BATCH_SIZE % device_count == 0, 'The batch size is not divisible by the device count' Log.i('Start training with {:d} GPUs ({:d} batches per GPU)'.format( torch.cuda.device_count(), Config.BATCH_SIZE // torch.cuda.device_count())) while not should_stop: for _, (_, image_batch, _, bboxes_batch, labels_batch) in enumerate(dataloader): #训练使用的数据集采用voc2007 batch_size = image_batch.shape[0] #(1,) image_batch = image_batch.cuda() #(1,3,h,w) bboxes_batch = bboxes_batch.cuda() #(1,gt_n,4) labels_batch = labels_batch.cuda() #(1,gt_n) anchor_objectness_losses, anchor_transformer_losses, proposal_class_losses, proposal_transformer_losses = \ model.train().forward(image_batch, bboxes_batch, labels_batch) #rpn的损失 anchor_objectness_loss = anchor_objectness_losses.mean() anchor_transformer_loss = anchor_transformer_losses.mean() #detection的损失 proposal_class_loss = proposal_class_losses.mean() proposal_transformer_loss = proposal_transformer_losses.mean() loss = anchor_objectness_loss + anchor_transformer_loss + proposal_class_loss + proposal_transformer_loss optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() losses.append(loss.item()) summary_writer.add_scalar('train/anchor_objectness_loss', anchor_objectness_loss.item(), step) summary_writer.add_scalar('train/anchor_transformer_loss', anchor_transformer_loss.item(), step) summary_writer.add_scalar('train/proposal_class_loss', proposal_class_loss.item(), step) summary_writer.add_scalar('train/proposal_transformer_loss', proposal_transformer_loss.item(), step) summary_writer.add_scalar('train/loss', loss.item(), step) step += 1 if step == num_steps_to_finish: should_stop = True if step % num_steps_to_display == 0: elapsed_time = time.time() - time_checkpoint time_checkpoint = time.time() steps_per_sec = num_steps_to_display / elapsed_time samples_per_sec = batch_size * steps_per_sec eta = (num_steps_to_finish - step) / steps_per_sec / 3600 avg_loss = sum(losses) / len(losses) lr = scheduler.get_lr()[0] #lr = optimizer.param_groups[0]['lr'] Log.i( f'[Step {step}] Avg. Loss = {avg_loss:.6f}, Learning Rate = {lr} ({samples_per_sec:.2f} samples/sec; ETA {eta:.1f} hrs)' ) #test if step == 10: path_to_checkpoint = model.module.save(path_to_checkpoints_dir, step, optimizer, scheduler) #path_to_checkpoint = model.module.save(path_to_checkpoints_dir, step, optimizer) Log.i(f'Model has been saved to {path_to_checkpoint}') if step % num_steps_to_snapshot == 0 or should_stop: path_to_checkpoint = model.module.save(path_to_checkpoints_dir, step, optimizer, scheduler) #path_to_checkpoint = model.module.save(path_to_checkpoints_dir, step, optimizer) Log.i(f'Model has been saved to {path_to_checkpoint}') if should_stop: break Log.i('Done')
def _train(dataset_name: str, backbone_name: str, path_to_data_dir: str, path_to_checkpoints_dir: str, path_to_resuming_checkpoint: Optional[str]): dataset = DatasetBase.from_name(dataset_name)(path_to_data_dir, DatasetBase.Mode.TRAIN, Config.IMAGE_MIN_SIDE, Config.IMAGE_MAX_SIDE) dataloader = DataLoader(dataset, batch_size=Config.BATCH_SIZE, sampler=DatasetBase.NearestRatioRandomSampler(dataset.image_ratios, num_neighbors=Config.BATCH_SIZE), num_workers=8, collate_fn=DatasetBase.padding_collate_fn, pin_memory=True) Log.i('Found {:d} samples'.format(len(dataset))) backbone = BackboneBase.from_name(backbone_name)(pretrained=True) model = nn.DataParallel( Model( backbone, dataset.num_classes(), pooler_mode=Config.POOLER_MODE, anchor_ratios=Config.ANCHOR_RATIOS, anchor_sizes=Config.ANCHOR_SIZES, rpn_pre_nms_top_n=Config.RPN_PRE_NMS_TOP_N, rpn_post_nms_top_n=Config.RPN_POST_NMS_TOP_N, anchor_smooth_l1_loss_beta=Config.ANCHOR_SMOOTH_L1_LOSS_BETA, proposal_smooth_l1_loss_beta=Config.PROPOSAL_SMOOTH_L1_LOSS_BETA ).cuda() ) optimizer = optim.SGD(model.parameters(), lr=Config.LEARNING_RATE, momentum=Config.MOMENTUM, weight_decay=Config.WEIGHT_DECAY) scheduler = WarmUpMultiStepLR(optimizer, milestones=Config.STEP_LR_SIZES, gamma=Config.STEP_LR_GAMMA, factor=Config.WARM_UP_FACTOR, num_iters=Config.WARM_UP_NUM_ITERS) step = 0 time_checkpoint = time.time() losses = deque(maxlen=100) summary_writer = SummaryWriter(os.path.join(path_to_checkpoints_dir, 'summaries')) should_stop = False num_steps_to_display = Config.NUM_STEPS_TO_DISPLAY num_steps_to_snapshot = Config.NUM_STEPS_TO_SNAPSHOT num_steps_to_finish = Config.NUM_STEPS_TO_FINISH if path_to_resuming_checkpoint is not None: step = model.module.load(path_to_resuming_checkpoint, optimizer, scheduler) Log.i(f'Model has been restored from file: {path_to_resuming_checkpoint}') device_count = torch.cuda.device_count() assert Config.BATCH_SIZE % device_count == 0, 'The batch size is not divisible by the device count' Log.i('Start training with {:d} GPUs ({:d} batches per GPU)'.format(torch.cuda.device_count(), Config.BATCH_SIZE // torch.cuda.device_count())) while not should_stop: for _, (_, image_batch, _, bboxes_batch, labels_batch) in enumerate(dataloader): batch_size = image_batch.shape[0] image_batch = image_batch.cuda() bboxes_batch = bboxes_batch.cuda() labels_batch = labels_batch.cuda() anchor_objectness_losses, anchor_transformer_losses, proposal_class_losses, proposal_transformer_losses = \ model.train().forward(image_batch, bboxes_batch, labels_batch) anchor_objectness_loss = anchor_objectness_losses.mean() anchor_transformer_loss = anchor_transformer_losses.mean() proposal_class_loss = proposal_class_losses.mean() proposal_transformer_loss = proposal_transformer_losses.mean() loss = anchor_objectness_loss + anchor_transformer_loss + proposal_class_loss + proposal_transformer_loss optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() losses.append(loss.item()) summary_writer.add_scalar('train/anchor_objectness_loss', anchor_objectness_loss.item(), step) summary_writer.add_scalar('train/anchor_transformer_loss', anchor_transformer_loss.item(), step) summary_writer.add_scalar('train/proposal_class_loss', proposal_class_loss.item(), step) summary_writer.add_scalar('train/proposal_transformer_loss', proposal_transformer_loss.item(), step) summary_writer.add_scalar('train/loss', loss.item(), step) step += 1 if step == num_steps_to_finish: should_stop = True if step % num_steps_to_display == 0: elapsed_time = time.time() - time_checkpoint time_checkpoint = time.time() steps_per_sec = num_steps_to_display / elapsed_time samples_per_sec = batch_size * steps_per_sec eta = (num_steps_to_finish - step) / steps_per_sec / 3600 avg_loss = sum(losses) / len(losses) lr = scheduler.get_lr()[0] Log.i(f'[Step {step}] Avg. Loss = {avg_loss:.6f}, Learning Rate = {lr:.8f} ({samples_per_sec:.2f} samples/sec; ETA {eta:.1f} hrs)') if step % num_steps_to_snapshot == 0 or should_stop: path_to_checkpoint = model.module.save(path_to_checkpoints_dir, step, optimizer, scheduler) Log.i(f'Model has been saved to {path_to_checkpoint}') if should_stop: break Log.i('Done')
# If specified we start from checkpoint # if config.TRAIN.PRETRAINED_WEIGHTS is not None and config.TRAIN.PRETRAINED_WEIGHTS != '': # if config.TRAIN.PRETRAINED_WEIGHTS.endswith(".pth"): # backbone.load_state_dict(torch.load(config.TRAIN.PRETRAINED_WEIGHTS)) # else: # backbone.load_darknet_weights(config.TRAIN.PRETRAINED_WEIGHTS) # Get dataloader dataset = DatasetBase.from_name('tiny-person')( config.TRAIN.PATH_TO_IMAGES_DIR, config.TRAIN.PATH_TO_ANNOTATIONS, DatasetBase.Mode.TRAIN) dataloader = DataLoader(dataset, batch_size=config.TRAIN.BATCH_SIZE, sampler=DatasetBase.NearestRatioRandomSampler( dataset.image_ratios, num_neighbors=config.TRAIN.BATCH_SIZE), num_workers=config.TRAIN.NUM_WORKERS, collate_fn=dataset.collate_fn, pin_memory=True) optimizer = torch.optim.Adam(model.parameters()) metrics = [ "grid_size", "loss", "x", "y", "w", "h", "conf", "cls",