def main(args): create_time = time.strftime('%Y%m%d_%H%M', time.localtime(time.time())) save_folder_path = os.path.join(args.save_folder, create_time) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # dataset = customDetection(root = args.image_root, # json_path = args.annotation, # transform = BaseTransform(img_size = args.image_size), # target_transform = customAnnotationTransform()) dataset = COCODetection(root=args.image_root, annotation_json=args.annotation, transform=BaseTransform(img_size=args.image_size), target_transform=COCOAnnotationTransform) dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, collate_fn=detection_collate) n_classes = dataset.get_class_number() + 1 print("Detect class number: {}".format(n_classes)) ## write category id to label name map dataset.get_class_map() model = mobilenetv3(n_classes=n_classes) ssd = ssd_mobilenetv3(model, n_classes) if args.pretrain_model_path: ssd.load_state_dict(torch.load(args.pretrain_model_path)) # Initialize the optimizer, with twice the default learning rate for biases, as in the original Caffe repo biases = list() not_biases = list() for param_name, param in model.named_parameters(): if param.requires_grad: if param_name.endswith('.bias'): biases.append(param) else: not_biases.append(param) optimizer = torch.optim.SGD(params=[{ 'params': biases, 'lr': args.learning_rate }, { 'params': not_biases }], lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) ssd = ssd.to(device) criterion = MultiBoxLossV3(ssd.priors_cxcy, args.threshold, args.neg_pos_ratio).to(device) print(f"epochs: {args.epochs}") for param_group in optimizer.param_groups: optimizer.param_groups[1]['lr'] = args.learning_rate print(f"learning rate. The new LR is {optimizer.param_groups[1]['lr']}") scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=15, verbose=True, threshold=0.00001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-08) n_train = min(dataset.__len__(), 5000) global_step = 0 writer = SummaryWriter() for epoch in range(args.epochs): mean_loss = 0 inference_count = 0 ssd.train() mean_count = 0 with tqdm(total=n_train, desc=f"{epoch + 1} / {args.epochs}", unit='img') as pbar: for img, target in dataloader: img = img.to(device) # target = [anno.to(device) for anno in target] # print(target) # boxes = target[:, :-1] # labels = target[:, -1] boxes = [anno.to(device)[:, :-1] for anno in target] labels = [anno.to(device)[:, -1] for anno in target] prediction_location_loss, prediction_confidence_loss = ssd(img) loss = criterion(prediction_location_loss, prediction_confidence_loss, boxes, labels) pbar.set_postfix(**{"loss ": float(loss)}) mean_loss += float(loss) mean_count += 1 optimizer.zero_grad() loss.backward() optimizer.step() pbar.update(img.shape[0]) scheduler.step(mean_loss) writer.add_scalar('Train/Loss', float(mean_loss / mean_count), global_step) global_step += 1 if epoch % 10 == 0 or epoch == args.epochs - 1: save_model(save_folder_path, ssd, epoch) writer.close()
def train(args): create_time = time.strftime('%Y%m%d_%H%M', time.localtime(time.time())) save_folder_path = os.path.join(args.save_folder, create_time) # n_classes = [20, 80][args.dataset == 'COCO'] # n_classes = 91 if not ((args.train_image_folder and args.val_image_folder) or args.annotation): print("train/val image folder and annotation should not be None") return train_dataset = COCODetection( root=args.root, image_set=args.train_image_folder, annotation_json=args.annotation, transform=SSDAugmentation(img_size=args.image_size), # transform = BaseTransform(img_size = args.image_size), target_transform=COCOAnnotationTransform()) train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=detection_collate) val_dataset = COCODetection( root=args.root, image_set=args.val_image_folder, annotation_json=args.annotation, transform=BaseTransform(img_size=args.image_size), target_transform=COCOAnnotationTransform()) n_classes = train_dataset.get_class_size() + 1 if args.class_map_path: train_dataset.get_class_map(args.class_map_path) if args.model == "mobilenetv2": model = MobileNetv2( n_classes=n_classes, width_mult=args.width_mult, round_nearest=8, dropout_ratio=args.dropout_ratio, use_batch_norm=True, ) ssd = create_mobilenetv2_ssd_lite(model, n_classes, width_mult=args.width_mult, use_batch_norm=True) elif args.model == "mobilenetv3": model = MobileNetv3(model_mode=args.model_mode, n_classes=n_classes, width_mult=args.width_mult, dropout_ratio=args.dropout_ratio) ssd = create_mobilenetv3_ssd_lite(model, n_classes, model_mode=args.model_mode) else: print("model structure only accept mobilenetv2 or mobilenetv3") return print("builded ssd module") if GPU: import torch.backends.cudnn as cudnn model.cuda() ssd.cuda() cudnn.benchmark = True if args.pretrain_model: ssd.load_state_dict( torch.load(args.pretrain_model, map_location=torch.device('cpu'))) elif args.pretrain_tfmodel and args.pretrain_tfmodel_weight_list: ssd_state_dict = ssd.state_dict() tf_weights_dict = load_tf_weights(args, ssd_state_dict) ssd.load_state_dict(tf_weights_dict) optimizer = optim.Adam(ssd.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay) criterion = MultiBoxLoss(n_classes, overlap_thresh=args.overlap_threshold, prior_for_matching=True, bkg_label=0, neg_mining=True, neg_pos=args.neg_pos_ratio, neg_overlap=0.5, encode_target=False) with torch.no_grad(): if args.model == "mobilenetv2": prior_box = PriorBox(MOBILEV2_300) elif args.model == "mobilenetv3": prior_box = PriorBox(MOBILEV3_300) priors = Variable(prior_box.forward()) print("created default bbox") n_train = min(train_dataset.__len__(), 5000) n_val = min(val_dataset.__len__(), 1000) global_step = 0 val_global_step = 0 writer = SummaryWriter(log_dir=args.summary_path) for epoch in range(args.epochs): mean_loss_conf = 0 mean_loss_loc = 0 inference_count = 0 ssd.train() with tqdm(total=n_train, desc=f"{epoch + 1} / {args.epochs}", unit='img') as pbar: for img, target in train_dataloader: if GPU: img = Variable(img.cuda()) target = [Variable(anno.cuda()) for anno in target] else: img = Variable(img) target = [Variable(anno) for anno in target] optimizer.zero_grad() inference = ssd(img) loss_loc, loss_conf = criterion(inference, priors, target) writer.add_scalar('Train/location_loss', float(loss_loc), global_step) writer.add_scalar('Train/confidence_loss', float(loss_conf), global_step) pbar.set_postfix( **{ "location loss": float(loss_loc), "confidence loss": float(loss_conf) }) mean_loss_loc += float(loss_loc) mean_loss_conf += float(loss_conf) total_loss = loss_loc + loss_conf total_loss.backward() # # clip gradient # # clip_grad_norm_(net.parameters(), 0.1) optimizer.step() pbar.update(img.shape[0]) global_step += 1 inference_count += img.shape[0] if inference_count > n_train: break pbar.set_postfix( **{ "location loss": float(mean_loss_loc / n_train), "confidence loss": float(mean_loss_conf / n_train) }) ssd.eval() val_mean_loss_loc = 0 val_mean_loss_conf = 0 with tqdm(total=n_val, desc="Validation", unit="img") as vpbar: for i in range(n_val): img = val_dataset.get_image(i) img = cv2.resize(img, (args.image_size, args.image_size)) height, width, _ = img.shape target = val_dataset.get_annotation(i, width, height) if GPU: img = torch.from_numpy( np.expand_dims(img.transpose(2, 0, 1), 0)).to(dtype=torch.float32).cuda() target = torch.FloatTensor(target).unsqueeze(0).cuda() else: img = torch.from_numpy( np.expand_dims(img.transpose(2, 0, 1), 0)).to(dtype=torch.float32) target = torch.FloatTensor(target).unsqueeze(0) inference = ssd(img) loss_loc, loss_conf = criterion(inference, priors, target) val_mean_loss_loc += float(loss_loc) val_mean_loss_conf += float(loss_conf) vpbar.set_postfix( **{ 'location loss': float(loss_loc), 'confidnece loss': float(loss_conf) }) vpbar.update(1) vpbar.set_postfix( **{ 'location loss': float(val_mean_loss_loc / n_val), 'confidnece loss': float(val_mean_loss_conf / n_val) }) writer.add_scalar('Test/location_loss', float(val_mean_loss_loc / n_val), val_global_step) writer.add_scalar('Test/confidence_loss', float(val_mean_loss_conf / n_val), val_global_step) val_global_step += 1 if epoch % 10 == 0 or epoch == args.epochs - 1: save_model(save_folder_path, ssd, epoch) writer.close()