def train(self, epochs=100, init_epoch=0):
# Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/
from torch.utils.tensorboard import SummaryWriter
self.tb_writer = SummaryWriter(comment=self.save_trial_id[:3])
for epoch_num in range(init_epoch + 1, epochs + 1):
print('total epochs: ', epochs)
self.retinanet.train()
self.retinanet.freeze_bn()
epoch_loss = []
loss_hist = collections.deque(maxlen=500)
total_num_iterations = len(self.dataloader_train)
dataloader_iterator = iter(self.dataloader_train)
pbar = tqdm(total=total_num_iterations)
for iter_num in range(1, total_num_iterations + 1):
try:
data = next(dataloader_iterator)
self.optimizer.zero_grad()
classification_loss, regression_loss = self.retinanet([
data['img'].to(device=self.device).float(),
data['annot'].to(device=self.device)
])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(self.retinanet.parameters(),
0.1)
self.optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
s = 'Trial {} -- Epoch: {}/{} | Iteration: {}/{} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(
self.save_trial_id[:3], epoch_num, epochs, iter_num,
total_num_iterations, float(classification_loss),
float(regression_loss), np.mean(loss_hist))
pbar.set_description(s)
pbar.update()
del classification_loss
del regression_loss
except Exception as e:
logger.info(e)
pbar.update()
continue
pbar.close()
self.scheduler.step(np.mean(epoch_loss))
self.final_epoch = epoch_num == epochs
mAP = csv_eval.evaluate(self.dataset_val, self.retinanet)
self._write_to_tensorboard(mAP, np.mean(loss_hist), epoch_num)
self._save_checkpoint(mAP, epoch_num)
if self.final_epoch:
self._save_classes_for_inference()
# last epoch delete last checkpoint and leave the best checkpoint
os.remove(self.save_last_checkpoint_path)
def main():
args = parse_args()
assert args.out or args.show or args.json_out, \
('Please specify at least one operation (save or show the results) '
'with the argument "--out" or "--show" or "--json_out"')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
if args.json_out is not None and args.json_out.endswith('.json'):
args.json_out = args.json_out[:-5]
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
# load checkpoint
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs, annotations = single_gpu_test(model, data_loader, args.show)
else:
model = MMDistributedDataParallel(model.cuda())
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
if args.eval:
_ = csv_eval.evaluate(outputs, annotations,
args.checkpoint) # 这里传入的似乎是dataset
def main(args=None):
parser = argparse.ArgumentParser(description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset', help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument('--csv_classes', help='Path to file containing class list (see readme)')
parser.add_argument('--csv_val', help='Path to file containing validation annotations (optional, see readme)')
parser.add_argument('--score_threshold', help='Score above which boxes are kept',default=0.5)
parser.add_argument('--nms_threshold', help='Score above which boxes are kept',default=0.5)
parser.add_argument('--model', help='Path to model (.pt) file.')
parser = parser.parse_args(args)
if parser.dataset == 'coco':
dataset_val = CocoDataset(parser.coco_path, set_name='val2017', transform=transforms.Compose([Normalizer(), Resizer()]))
elif parser.dataset == 'csv':
dataset_val = CSVDataset(train_file=parser.csv_val, class_list=parser.csv_classes, transform=transforms.Compose([Normalizer(), Resizer()]))
else:
raise ValueError('Dataset type not understood (must be csv or coco), exiting.')
sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
dataloader_val = DataLoader(dataset_val, num_workers=0, collate_fn=collater, batch_sampler=sampler_val,shuffle=False)
retinanet = torch.load(parser.model)
score_threshold = float(parser.score_threshold)
nms_threshold = float(parser.score_threshold)
use_gpu = True
f = open('mAPs.txt','w')
writer = csv.writer(f,delimiter = ",")
if use_gpu:
retinanet = retinanet.cuda()
retinanet.eval()
thresholds = np.array([0.1 + 0.05*n for n in range(10)])
for nms_threshold in thresholds:
for score_threshold in thresholds:
mAPs = csv_eval.evaluate(dataset_val, retinanet,iou_threshold=0.5,score_threshold=score_threshold,nms_threshold = nms_threshold)
maps=np.mean([ap[0] for ap in mAPs.values()])
writer.writerow([score_threshold,nms_threshold,maps])
def test(self, validation=False):
tqdm_loader = tqdm(self.test_loader, total=len(self.test_loader))
self.change_model_state('eval')
mAP = csv_eval.evaluate(self.dataset_test, self.test_loader,
self.model, self.log_file)
return
# array to save loss for testing data
all_detections = [[None for i in range(3)]
for j in range(len(self.test_loader))]
for step, batch in enumerate(tqdm_loader):
scores, labels, boxes = self.feed_into_net(batch, train=False)
scores = scores.cpu().data.numpy()
labels = labels.cpu().data.numpy()
boxes = boxes.cpu().data.numpy()
if validation is True:
anno = batch['annot']
indices = np.where(scores >= 0.00)[0]
if indices.shape[0] > 0:
scores = scores[indices]
scores_sort = np.argsort(-scores)[:100]
image_boxes = boxes[indices[scores_sort], :]
image_scores = scores[scores_sort]
image_labels = labels[indices[scores_sort]]
image_detections = np.concatenate([image_boxes, \
image_scores.reshape(-1, 1), \
image_labels.reshape(-1, 1)], axis=1)
for label in range(3):
all_detections[step][label] = image_detections[
image_detections[:, -1] == label, :-1]
exit()
from torch.utils.data import Dataset, DataLoader
from dataloader import CocoDataset, CSVDataset, collater, Resizer, AspectRatioBasedSampler, Augmenter, UnNormalizer, Normalizer, TXTDataset
from torchvision import datasets, models, transforms
import argparse
import csv_eval
ap = argparse.ArgumentParser(description='Just for eval.')
ap.add_argument('-B', '--ImageNamePath')
ap.add_argument('-A', '--AnnotationPath')
ap.add_argument('-I', '--ImagePath')
args = vars(ap.parse_args())
dataset_val = TXTDataset(ImgNamePth=args['ImageNamePath'],
AnnoPth=args['AnnotationPath'],
ImgPth=args['ImagePath'],
transform=transforms.Compose(
[Normalizer(), Resizer()]))
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler_val)
modelpath = os.path.join(os.path.dirname(__file__), 'weights/model.pt')
modelretina = torch.load(modelpath)
modelretina.cuda()
modelretina.eval()
AP = csv_eval.evaluate(dataset_val, modelretina)
map = (AP[0][0] + AP[1][0]) / 2
print("mAp:" + str(map))
def main(args=None):
parser = argparse.ArgumentParser(description='Training script for training a EfficientDet network.')
parser.add_argument('--dataset', help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument('--csv_train', help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes', help='Path to file containing class list (see readme)')
parser.add_argument('--csv_val', help='Path to file containing validation annotations (optional, see readme)')
parser.add_argument('--phi', help='EfficientNet scaling coefficient.', type=int, default=0)
parser.add_argument('--batch-size', help='Batch size', type=int, default=8)
parser.add_argument('--epochs', help='Number of epochs', type=int, default=100)
parser = parser.parse_args(args)
# Create the data loaders
if parser.dataset == 'coco':
if parser.coco_path is None:
raise ValueError('Must provide --coco_path when training on COCO,')
dataset_train = CocoDataset(parser.coco_path, set_name='train2017', transform=transforms.Compose([Normalizer(), Augmenter(), Resizer(img_size=512)]))
dataset_val = CocoDataset(parser.coco_path, set_name='val2017', transform=transforms.Compose([Normalizer(), Resizer(img_size=512)]))
elif parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO')
if parser.csv_classes is None:
raise ValueError('Must provide --csv_classes when training on COCO')
dataset_train = CSVDataset(train_file=parser.csv_train, class_list=parser.csv_classes, transform=transforms.Compose([Normalizer(), Augmenter(), Resizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val, class_list=parser.csv_classes, transform=transforms.Compose([Normalizer(), Resizer()]))
else:
raise ValueError('Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train, batch_size=parser.batch_size, drop_last=False)
dataloader_train = DataLoader(dataset_train, num_workers=3, collate_fn=collater, batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
dataloader_val = DataLoader(dataset_val, num_workers=3, collate_fn=collater, batch_sampler=sampler_val)
# Create the model
efficientdet = model.efficientdet(num_classes=dataset_train.num_classes(), pretrained=True, phi=parser.phi)
use_gpu = True
if use_gpu:
efficientdet = efficientdet.cuda()
efficientdet = torch.nn.DataParallel(efficientdet).cuda()
efficientdet.training = True
optimizer = optim.Adam(efficientdet.parameters(), lr=1e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True)
loss_hist = collections.deque(maxlen=500)
efficientdet.train()
efficientdet.module.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
for epoch_num in range(parser.epochs):
efficientdet.train()
efficientdet.module.freeze_bn()
epoch_loss = []
print(('\n' + '%10s' * 5) % ('Epoch', 'gpu_mem', 'Loss', 'cls', 'rls'))
pbar = tqdm(enumerate(dataloader_train), total=len(dataloader_train))
for iter_num, data in pbar:
try:
optimizer.zero_grad()
classification_loss, regression_loss = efficientdet([data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(efficientdet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
loss = (loss * iter_num) / (iter_num + 1) # update mean losses
mem = torch.cuda.memory_cached() / 1E9 if torch.cuda.is_available() else 0 # (GB)
s = ('%10s' * 2 + '%10.3g' * 3) % (
'%g/%g' % (epoch_num, parser.epochs - 1), '%.3gG' % mem, np.mean(loss_hist), float(regression_loss), float(classification_loss))
pbar.set_description(s)
del classification_loss
del regression_loss
except Exception as e:
raise(e)
continue
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, efficientdet)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, efficientdet)
scheduler.step(np.mean(epoch_loss))
torch.save(efficientdet.module, '{}_retinanet_{}.pt'.format(parser.dataset, epoch_num))
efficientdet.eval()
torch.save(efficientdet, 'model_final.pt'.format(epoch_num))
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
#add a bunch of arguments(customized by Yu Han Huang)
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument('--model', default='None')
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--resnext',
help='change backbone to resnext101',
action='store_true')
parser.add_argument('--epochs',
help='Number of Epochs',
type=int,
default=12)
parser.add_argument('--batch_size', help='Batch Size', type=int, default=4)
parser.add_argument('--workers',
help='Number of Workers',
type=int,
default=4)
parser.add_argument('--lr',
help='Learning Rate for training',
type=float,
default=1e-5)
parser.add_argument(
'--dropout1',
help='Dropout Rate for layer dropout1 in ClassficationModel',
type=float,
default=0.25)
parser.add_argument(
'--dropout2',
help='Dropout Rate for layer dropout2 in ClassficationModel',
type=float,
default=0.25)
parser.add_argument(
'--angle',
help='Angle of pictures while implementing Data Augmentation',
type=float,
default=6)
parser.add_argument('--size',
help='The length of the side of pictures',
type=int,
default=512)
parser.add_argument(
'--zoom_range',
help=
'Zoom Range of pictures while implementing Data Augmentation. Please type two arguments for this one.',
nargs='+',
type=float,
default=[-0.1, 0.1])
parser.add_argument('--alpha',
help='Alpha for focal loss',
type=float,
default=0.25)
parser.add_argument('--gamma',
help='Gamma for focal loss',
type=float,
default=2)
parser.add_argument('--loss_with_no_bboxes', action='store_true')
parser.add_argument('--no_bboxes_alpha',
help='Alpha for focal loss',
type=float,
default=0.5)
parser.add_argument('--no_bboxes_gamma',
help='Gamma for focal loss',
type=float,
default=2)
parser = parser.parse_args(args)
# Create the data loaders
if parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on CSV,')
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on CSV,')
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose([
Normalizer(),
Augmenter(angle=parser.angle),
Resizer(zoom_range=parser.zoom_range,
side=parser.side)
]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
ValResizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=parser.batch_size,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=parser.workers,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
# add arguments dropout1, dropout2, alpha, gamma, loss_with_no_bboxes, no_bboxes_alpha, no_bboxes_gamma(customized by Yu Han Huang)
if parser.resnext == False:
if parser.depth == 18:
retinanet = model.resnet18(
num_classes=dataset_train.num_classes(),
pretrained=True,
dropout1=parser.dropout1,
dropout2=parser.dropout2,
alpha=parser.alpha,
gamma=parser.gamma,
loss_with_no_bboxes=parser.loss_with_no_bboxes,
no_bboxes_alpha=parser.no_bboxes_alpha,
no_bboxes_gamma=parser.no_bboxes_gamma)
elif parser.depth == 34:
retinanet = model.resnet34(
num_classes=dataset_train.num_classes(),
pretrained=True,
dropout1=parser.dropout1,
dropout2=parser.dropout2,
alpha=parser.alpha,
gamma=parser.gamma,
loss_with_no_bboxes=parser.loss_with_no_bboxes,
no_bboxes_alpha=parser.no_bboxes_alpha,
no_bboxes_gamma=parser.no_bboxes_gamma)
elif parser.depth == 50:
retinanet = model.resnet50(
num_classes=dataset_train.num_classes(),
pretrained=True,
dropout1=parser.dropout1,
dropout2=parser.dropout2,
alpha=parser.alpha,
gamma=parser.gamma,
loss_with_no_bboxes=parser.loss_with_no_bboxes,
no_bboxes_alpha=parser.no_bboxes_alpha,
no_bboxes_gamma=parser.no_bboxes_gamma)
elif parser.depth == 101:
retinanet = model.resnet101(
num_classes=dataset_train.num_classes(),
pretrained=True,
dropout1=parser.dropout1,
dropout2=parser.dropout2,
alpha=parser.alpha,
gamma=parser.gamma,
loss_with_no_bboxes=parser.loss_with_no_bboxes,
no_bboxes_alpha=parser.no_bboxes_alpha,
no_bboxes_gamma=parser.no_bboxes_gamma)
elif parser.depth == 152:
retinanet = model.resnet152(
num_classes=dataset_train.num_classes(),
pretrained=True,
dropout1=parser.dropout1,
dropout2=parser.dropout2,
alpha=parser.alpha,
gamma=parser.gamma,
loss_with_no_bboxes=parser.loss_with_no_bboxes,
no_bboxes_alpha=parser.no_bboxes_alpha,
no_bboxes_gamma=parser.no_bboxes_gamma)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
else:
if parser.depth == 101:
retinanet = model.resnext101(
num_classes=dataset_train.num_classes(),
pretrained=True,
dropout1=parser.dropout1,
dropout2=parser.dropout2,
alpha=parser.alpha,
gamma=parser.gamma,
loss_with_no_bboxes=parser.loss_with_no_bboxes,
no_bboxes_alpha=parser.no_bboxes_alpha,
no_bboxes_gamma=parser.no_bboxes_gamma)
use_gpu = True
if parser.model != 'None':
retinanet = torch.load(parser.model)
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=parser.lr)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
for epoch_num in range(parser.epochs):
retinanet.train()
retinanet.module.freeze_bn()
print_activate = 0
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
try:
optimizer.zero_grad()
classification_loss, regression_loss = retinanet(
[data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
#print(classification_loss, regression_loss)
if bool(loss == 0):
continue
loss.backward()
print_activate += 1
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
if print_activate % 15 == 0:
print(
'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'
.format(epoch_num, iter_num,
float(classification_loss),
float(regression_loss), np.mean(loss_hist)))
del loss
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
scheduler.step(np.mean(epoch_loss))
torch.save(
retinanet.module,
'{}_retinanet_resnext_v4_{}.pt'.format(parser.dataset, epoch_num))
if parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet)
retinanet.eval()
torch.save(retinanet, 'model_final.pt'.format(epoch_num))
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.',
default="csv")
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)',
default="binary_class.csv")
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=18)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=500)
parser.add_argument('--epochs_only_det',
help='Number of epochs to train detection part',
type=int,
default=1)
parser.add_argument('--max_epochs_no_improvement',
help='Max epochs without improvement',
type=int,
default=100)
parser.add_argument('--pretrained_model',
help='Path of .pt file with pretrained model',
default='esposallescsv_retinanet_0.pt')
parser.add_argument('--model_out',
help='Path of .pt file with trained model to save',
default='trained')
parser.add_argument('--score_threshold',
help='Score above which boxes are kept',
type=float,
default=0.5)
parser.add_argument('--nms_threshold',
help='Score above which boxes are kept',
type=float,
default=0.2)
parser.add_argument('--max_boxes',
help='Max boxes to be fed to recognition',
default=95)
parser.add_argument('--seg_level',
help='[line, word], to choose anchor aspect ratio',
default='word')
parser.add_argument(
'--early_stop_crit',
help='Early stop criterion, detection (map) or transcription (cer)',
default='cer')
parser.add_argument('--max_iters_epoch',
help='Max steps per epoch (for debugging)',
default=1000000)
parser.add_argument('--train_htr',
help='Train recognition or not',
default='True')
parser.add_argument('--train_det',
help='Train detection or not',
default='True')
parser.add_argument(
'--binary_classifier',
help=
'Wether to use classification branch as binary or not, multiclass instead.',
default='False')
parser.add_argument(
'--htr_gt_box',
help='Train recognition branch with box gt (for debugging)',
default='False')
parser.add_argument(
'--ner_branch',
help='Train named entity recognition with separate branch',
default='False')
parser = parser.parse_args(args)
if parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train')
dataset_name = parser.csv_train.split("/")[-2]
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
# Files for training log
experiment_id = str(time.time()).split('.')[0]
valid_cer_f = open('trained_models/' + parser.model_out + 'log.txt', 'w')
for arg in vars(parser):
if getattr(parser, arg) is not None:
valid_cer_f.write(
str(arg) + ' ' + str(getattr(parser, arg)) + '\n')
current_commit = subprocess.check_output(['git', 'rev-parse', 'HEAD'])
valid_cer_f.write(str(current_commit))
valid_cer_f.write(
"epoch_num cer best cer mAP best mAP time\n")
valid_cer_f.close()
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=1,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=0,
collate_fn=collater,
batch_sampler=sampler_val)
if not os.path.exists('trained_models'):
os.mkdir('trained_models')
# Create the model
train_htr = parser.train_htr == 'True'
htr_gt_box = parser.htr_gt_box == 'True'
ner_branch = parser.ner_branch == 'True'
binary_classifier = parser.binary_classifier == 'True'
torch.backends.cudnn.benchmark = False
alphabet = dataset_train.alphabet
if os.path.exists(parser.pretrained_model):
retinanet = torch.load(parser.pretrained_model)
retinanet.classificationModel = ClassificationModel(
num_features_in=256,
num_anchors=retinanet.anchors.num_anchors,
num_classes=dataset_train.num_classes())
if ner_branch:
retinanet.nerModel = NERModel(
feature_size=256,
pool_h=retinanet.pool_h,
n_classes=dataset_train.num_classes(),
pool_w=retinanet.pool_w)
else:
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(),
pretrained=True,
max_boxes=int(parser.max_boxes),
score_threshold=float(
parser.score_threshold),
seg_level=parser.seg_level,
alphabet=alphabet,
train_htr=train_htr,
htr_gt_box=htr_gt_box,
ner_branch=ner_branch,
binary_classifier=binary_classifier)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(),
pretrained=True,
max_boxes=int(parser.max_boxes),
score_threshold=float(
parser.score_threshold),
seg_level=parser.seg_level,
alphabet=alphabet,
train_htr=train_htr,
htr_gt_box=htr_gt_box)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(
num_classes=dataset_train.num_classes(), pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(
num_classes=dataset_train.num_classes(), pretrained=True)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
train_htr = parser.train_htr == 'True'
train_det = parser.train_det == 'True'
retinanet.htr_gt_box = parser.htr_gt_box == 'True'
retinanet.train_htr = train_htr
retinanet.epochs_only_det = parser.epochs_only_det
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=50,
verbose=True)
loss_hist = collections.deque(maxlen=500)
ctc = CTCLoss()
retinanet.train()
retinanet.module.freeze_bn()
best_cer = 1000
best_map = 0
epochs_no_improvement = 0
verbose_each = 20
optimize_each = 1
objective = 100
best_objective = 10000
print(('Num training images: {}'.format(len(dataset_train))))
for epoch_num in range(parser.epochs):
cers = []
retinanet.training = True
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
if iter_num > int(parser.max_iters_epoch): break
try:
if iter_num % optimize_each == 0:
optimizer.zero_grad()
(classification_loss, regression_loss, ctc_loss,
ner_loss) = retinanet([
data['img'].cuda().float(), data['annot'], ctc, epoch_num
])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
if train_det:
if train_htr:
loss = ctc_loss + classification_loss + regression_loss + ner_loss
else:
loss = classification_loss + regression_loss + ner_loss
elif train_htr:
loss = ctc_loss
else:
continue
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
if iter_num % verbose_each == 0:
print((
'Epoch: {} | Step: {} |Classification loss: {:1.5f} | Regression loss: {:1.5f} | CTC loss: {:1.5f} | NER loss: {:1.5f} | Running loss: {:1.5f} | Total loss: {:1.5f}\r'
.format(epoch_num, iter_num,
float(classification_loss),
float(regression_loss), float(ctc_loss),
float(ner_loss), np.mean(loss_hist),
float(loss), "\r")))
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
torch.cuda.empty_cache()
except Exception as e:
print(e)
continue
if parser.dataset == 'csv' and parser.csv_val is not None and train_det:
print('Evaluating dataset')
mAP, text_mAP, current_cer = csv_eval.evaluate(
dataset_val, retinanet, score_threshold=parser.score_threshold)
#text_mAP,_ = csv_eval_binary_map.evaluate(dataset_val, retinanet,score_threshold=parser.score_threshold)
objective = current_cer * (1 - mAP)
retinanet.eval()
retinanet.training = False
retinanet.score_threshold = float(parser.score_threshold)
'''for idx,data in enumerate(dataloader_val):
if idx>int(parser.max_iters_epoch): break
print("Eval CER on validation set:",idx,"/",len(dataset_val),"\r")
image_name = dataset_val.image_names[idx].split('/')[-1].split('.')[-2]
#generate_pagexml(image_name,data,retinanet,parser.score_threshold,parser.nms_threshold,dataset_val)
text_gt =".".join(dataset_val.image_names[idx].split('.')[:-1])+'.txt'
f =open(text_gt,'r')
text_gt_lines=f.readlines()[0]
transcript_pred = get_transcript(image_name,data,retinanet,float(parser.score_threshold),float(parser.nms_threshold),dataset_val,alphabet)
cers.append(float(editdistance.eval(transcript_pred,text_gt_lines))/len(text_gt_lines))'''
t = str(time.time()).split('.')[0]
valid_cer_f.close()
#print("GT",text_gt_lines)
#print("PREDS SAMPLE:",transcript_pred)
if parser.early_stop_crit == 'cer':
if float(objective) < float(
best_objective): #float(current_cer)<float(best_cer):
best_cer = current_cer
best_objective = objective
epochs_no_improvement = 0
torch.save(
retinanet.module, 'trained_models/' + parser.model_out +
'{}_retinanet.pt'.format(parser.dataset))
else:
epochs_no_improvement += 1
if mAP > best_map:
best_map = mAP
elif parser.early_stop_crit == 'map':
if mAP > best_map:
best_map = mAP
epochs_no_improvement = 0
torch.save(
retinanet.module, 'trained_models/' + parser.model_out +
'{}_retinanet.pt'.format(parser.dataset))
else:
epochs_no_improvement += 1
if float(current_cer) < float(best_cer):
best_cer = current_cer
if train_det:
print(epoch_num, "mAP: ", mAP, " best mAP", best_map)
if train_htr:
print("VALID CER:", current_cer, "best CER", best_cer)
print("Epochs no improvement:", epochs_no_improvement)
valid_cer_f = open('trained_models/' + parser.model_out + 'log.txt',
'a')
valid_cer_f.write(
str(epoch_num) + " " + str(current_cer) + " " + str(best_cer) +
' ' + str(mAP) + ' ' + str(best_map) + ' ' + str(text_mAP) + '\n')
if epochs_no_improvement > 3:
for param_group in optimizer.param_groups:
if param_group['lr'] > 10e-5:
param_group['lr'] *= 0.1
if epochs_no_improvement >= parser.max_epochs_no_improvement:
print("TRAINING FINISHED AT EPOCH", epoch_num, ".")
sys.exit()
scheduler.step(np.mean(epoch_loss))
torch.cuda.empty_cache()
retinanet.eval()
def main(args=None):
parser = argparse.ArgumentParser(description='Simple testing script for RetinaNet network.')
parser.add_argument('--dataset', help='Dataset type, must be one of csv or coco.',default = "csv")
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument('--csv_classes', help='Path to file containing class list (see readme)',default="binary_class.csv")
parser.add_argument('--csv_val', help='Path to file containing validation annotations (optional, see readme)')
parser.add_argument('--csv_box_annot', help='Path to file containing predicted box annotations ')
parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152', type=int, default=18)
parser.add_argument('--epochs', help='Number of epochs', type=int, default=500)
parser.add_argument('--model', help='Path of .pt file with trained model',default = 'esposallescsv_retinanet_0.pt')
parser.add_argument('--model_out', help='Path of .pt file with trained model to save',default = 'trained')
parser.add_argument('--score_threshold', help='Score above which boxes are kept',default=0.15)
parser.add_argument('--nms_threshold', help='Score above which boxes are kept',default=0.2)
parser.add_argument('--max_epochs_no_improvement', help='Max epochs without improvement',default=100)
parser.add_argument('--max_boxes', help='Max boxes to be fed to recognition',default=50)
parser.add_argument('--seg_level', help='Line or word, to choose anchor aspect ratio',default='line')
parser.add_argument('--htr_gt_box',help='Train recognition branch with box gt (for debugging)',default=False)
parser = parser.parse_args(args)
# Create the data loaders
if parser.dataset == 'csv':
if parser.csv_classes is None:
raise ValueError('Must provide --csv_classes when training on COCO,')
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val, class_list=parser.csv_classes, transform=transforms.Compose([Normalizer(), Resizer()]))
if parser.csv_box_annot is not None:
box_annot_data = CSVDataset(train_file=parser.csv_box_annot, class_list=parser.csv_classes, transform=transforms.Compose([Normalizer(), Resizer()]))
else:
box_annot_data = None
else:
raise ValueError('Dataset type not understood (must be csv or coco), exiting.')
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
dataloader_val = DataLoader(dataset_val, num_workers=0, collate_fn=collater, batch_sampler=sampler_val)
if box_annot_data is not None:
sampler_val = AspectRatioBasedSampler(box_annot_data, batch_size=1, drop_last=False)
dataloader_box_annot = DataLoader(box_annot_data, num_workers=0, collate_fn=collater, batch_sampler=sampler_val)
else:
dataloader_box_annot = dataloader_val
if not os.path.exists('trained_models'):
os.mkdir('trained_models')
# Create the model
alphabet=dataset_val.alphabet
if os.path.exists(parser.model):
retinanet = torch.load(parser.model)
else:
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_val.num_classes(), pretrained=True,max_boxes=int(parser.max_boxes),score_threshold=float(parser.score_threshold),seg_level=parser.seg_level,alphabet=alphabet)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(), pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(), pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(), pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(), pretrained=True)
else:
raise ValueError('Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
#retinanet = torch.load('../Documents/TRAINED_MODELS/pytorch-retinanet/esposallescsv_retinanet_99.pt')
#print "LOADED pretrained MODEL\n\n"
optimizer = optim.Adam(retinanet.parameters(), lr=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=4, verbose=True)
loss_hist = collections.deque(maxlen=500)
ctc = CTCLoss()
retinanet.module.freeze_bn()
best_cer = 1000
epochs_no_improvement=0
cers=[]
retinanet.eval()
retinanet.module.epochs_only_det = 0
#retinanet.module.htr_gt_box = False
retinanet.training=False
if parser.score_threshold is not None:
retinanet.module.score_threshold = float(parser.score_threshold)
'''if parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
'''
mAP = csv_eval.evaluate(dataset_val, retinanet,score_threshold=retinanet.module.score_threshold)
aps = []
for k,v in mAP.items():
aps.append(v[0])
print ("VALID mAP:",np.mean(aps))
print("score th",retinanet.module.score_threshold)
for idx,data in enumerate(dataloader_box_annot):
print("Eval CER on validation set:",idx,"/",len(dataloader_box_annot),"\r")
if box_annot_data:
image_name = box_annot_data.image_names[idx].split('/')[-1].split('.')[-2]
else:
image_name = dataset_val.image_names[idx].split('/')[-1].split('.')[-2]
#generate_pagexml(image_name,data,retinanet,parser.score_threshold,parser.nms_threshold,dataset_val)
text_gt_path="/".join(dataset_val.image_names[idx].split('/')[:-1])
text_gt = os.path.join(text_gt_path,image_name+'.txt')
f =open(text_gt,'r')
text_gt_lines=f.readlines()[0]
transcript_pred = get_transcript(image_name,data,retinanet,retinanet.module.score_threshold,float(parser.nms_threshold),dataset_val,alphabet)
cers.append(float(editdistance.eval(transcript_pred,text_gt_lines))/len(text_gt_lines))
print("GT",text_gt_lines)
print("PREDS SAMPLE:",transcript_pred)
print("VALID CER:",np.mean(cers),"best CER",best_cer)
print("GT",text_gt_lines)
print("PREDS SAMPLE:",transcript_pred)
print("VALID CER:",np.mean(cers),"best CER",best_cer)
def main(args=None):
#def main(epoch):
parser = argparse.ArgumentParser(description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset', help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument('--csv_train', help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes', help='Path to file containing class list (see readme)')
parser.add_argument('--csv_val', help='Path to file containing validation annotations (optional, see readme)')
parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152', type=int, default=50)
parser.add_argument('--epochs', help='Number of epochs', type=int, default=100)
#parser.add_argument('--resume', '-r', action='store_true', help='resume from checkpoint')
parser.add_argument('--start-epoch', default=0, type=int, help='manual epoch number (useful on restarts)')
parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')
parser = parser.parse_args(args)
#args = parser.parse_args()
#parser = parser.parse_args(epoch)
# Create the data loaders
if parser.dataset == 'coco':
if parser.coco_path is None:
raise ValueError('Must provide --coco_path when training on COCO,')
dataset_train = CocoDataset(parser.coco_path, set_name='train2017', transform=transforms.Compose([Normalizer(), Augmenter(), Resizer()]))
dataset_val = CocoDataset(parser.coco_path, set_name='val2017', transform=transforms.Compose([Normalizer(), Resizer()]))
elif parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes is None:
raise ValueError('Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(train_file=parser.csv_train, class_list=parser.csv_classes, transform=transforms.Compose([Normalizer(), Augmenter(), Resizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val, class_list=parser.csv_classes, transform=transforms.Compose([Normalizer(), Resizer()]))
else:
raise ValueError('Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train, batch_size=4, drop_last=False)
dataloader_train = DataLoader(dataset_train, num_workers=3, collate_fn=collater, batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
dataloader_val = DataLoader(dataset_val, num_workers=3, collate_fn=collater, batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(), pretrained=True)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(), pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(), pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(), pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(), pretrained=True)
else:
raise ValueError('Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
#retinanet().load_state_dict(torch.load('/users/wenchi/ghwwc/Pytorch-retinanet-master/resnet50-19c8e357.pth'))
#if True:
#print('==> Resuming from checkpoint..')
#checkpoint = torch.load('/users/wenchi/ghwwc/Pytorch-retinanet-master/coco_retinanet_2.pt')
#retinanet().load_state_dict(checkpoint)
#best_loss = checkpoint['loss']
#start_epoch = checkpoint['epoch']
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
#optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)
optimizer = optim.SGD(retinanet.parameters(), lr=1e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
#retinanet.freeze_bn() #for train from a middle state
retinanet.module.freeze_bn() #for train from the very beginning
print('Num training images: {}'.format(len(dataset_train)))
for epoch_num in range(parser.start_epoch, parser.epochs):
if parser.resume:
if os.path.isfile(parser.resume):
print("=>loading checkpoint '{}'".format(parser.resume))
checkpoint = torch.load(parser.resume)
print(parser.start_epoch)
#parser.start_epoch = checkpoint['epoch']
#retinanet.load_state_dict(checkpoint['state_dict'])
retinanet=checkpoint
#retinanet.load_state_dict(checkpoint)
print(retinanet)
#optimizer.load_state_dict(checkpoint)
print("=> loaded checkpoint '{}' (epoch {})".format(parser.resume, checkpoint))
else:
print("=> no checkpoint found at '{}'".format(parser.resume))
retinanet.train()
retinanet.freeze_bn()
#retinanet.module.freeze_bn()
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, retinanet)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet)
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
try:
optimizer.zero_grad()
classification_loss, regression_loss = retinanet([data['img'].cuda().float(), data['annot'].cuda()])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
print('Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, retinanet)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet)
scheduler.step(np.mean(epoch_loss))
#torch.save(retinanet.module, '{}_retinanet_101_{}.pt'.format(parser.dataset, epoch_num))
torch.save(retinanet, '{}_retinanet_dilation_experiment1_{}.pt'.format(parser.dataset, epoch_num))
name = '{}_retinanet_dilation_experiment1_{}.pt'.format(parser.dataset, epoch_num)
parser.resume = '/users/wenchi/ghwwc/pytorch-retinanet-master_new/name'
retinanet.eval()
torch.save(retinanet, 'model_final_dilation_experiment1.pt'.format(epoch_num))
def get_best_metrics(self):
checkpoint = torch.load(self.save_best_checkpoint_path)
self.retinanet.load_state_dict(checkpoint['model'])
mAP = csv_eval.evaluate(self.dataset_val, self.retinanet)
return mAP.item()
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=100)
parser.add_argument('--optimizer',
help='[SGD | Adam]',
type=str,
default='SGD')
parser.add_argument('--model', help='Path to model (.pt) file.')
parser = parser.parse_args(args)
# Create the data loaders
print("\n[Phase 1]: Creating DataLoader for {} dataset".format(
parser.dataset))
if parser.dataset == 'coco':
if parser.coco_path is None:
raise ValueError('Must provide --coco_path when training on COCO,')
dataset_train = CocoDataset(parser.coco_path,
set_name='train2014',
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
dataset_val = CocoDataset(parser.coco_path,
set_name='val2014',
transform=transforms.Compose(
[Normalizer(), Resizer()]))
elif parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=8,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=8,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=16,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=8,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(),
pretrained=True)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
print('| Num training images: {}'.format(len(dataset_train)))
print('| Num test images : {}'.format(len(dataset_val)))
print("\n[Phase 2]: Preparing RetinaNet Detection Model...")
use_gpu = torch.cuda.is_available()
if use_gpu:
device = torch.device('cuda')
retinanet = retinanet.to(device)
retinanet = torch.nn.DataParallel(retinanet,
device_ids=range(
torch.cuda.device_count()))
print("| Using %d GPUs for Train/Validation!" % torch.cuda.device_count())
retinanet.training = True
if parser.optimizer == 'Adam':
optimizer = optim.Adam(retinanet.parameters(),
lr=1e-5) # not mentioned
print("| Adam Optimizer with Learning Rate = {}".format(1e-5))
elif parser.optimizer == 'SGD':
optimizer = optim.SGD(retinanet.parameters(),
lr=1e-2,
momentum=0.9,
weight_decay=1e-4)
print("| SGD Optimizer with Learning Rate = {}".format(1e-2))
else:
raise ValueError('Unsupported Optimizer, must be one of [SGD | Adam]')
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn(
) # Freeze the BN parameters to ImageNet configuration
# Check if there is a 'checkpoints' path
if not osp.exists('./checkpoints/'):
os.makedirs('./checkpoints/')
print("\n[Phase 3]: Training Model on {} dataset...".format(
parser.dataset))
for epoch_num in range(parser.epochs):
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
try:
optimizer.zero_grad()
classification_loss, regression_loss = retinanet(
[data['img'].to(device), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.001)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
sys.stdout.write('\r')
sys.stdout.write(
'| Epoch: {} | Iteration: {}/{} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'
.format(epoch_num + 1, iter_num + 1, len(dataloader_train),
float(classification_loss), float(regression_loss),
np.mean(loss_hist)))
sys.stdout.flush()
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
print("\n| Saving current best model at epoch {}...".format(epoch_num +
1))
torch.save(
retinanet.state_dict(),
'./checkpoints/{}_retinanet_{}.pt'.format(parser.dataset,
epoch_num + 1))
if parser.dataset == 'coco':
#print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, retinanet, device)
elif parser.dataset == 'csv' and parser.csv_val is not None:
#print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet, device)
scheduler.step(np.mean(epoch_loss))
retinanet.eval()
torch.save(retinanet.state_dict(), './checkpoints/model_final.pt')
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',
default="csv",
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument(
'--csv_train',
default="./data/train_only.csv",
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
default="./data/classes.csv",
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
default="./data/train_only.csv",
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument('--voc_train',
default="./data/voc_train",
help='Path to containing images and annAnnotations')
parser.add_argument('--voc_val',
default="./data/bov_train",
help='Path to containing images and annAnnotations')
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=101)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=40)
parser = parser.parse_args(args)
# Create the data loaders
if parser.dataset == 'coco':
if parser.coco_path is None:
raise ValueError('Must provide --coco_path when training on COCO,')
dataset_train = CocoDataset(parser.coco_path,
set_name='train2017',
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
dataset_val = CocoDataset(parser.coco_path,
set_name='val2017',
transform=transforms.Compose(
[Normalizer(), Resizer()]))
elif parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Resizer()]))
elif parser.dataset == 'voc':
if parser.voc_train is None:
raise ValueError(
'Must provide --voc_train when training on PASCAL VOC,')
dataset_train = XML_VOCDataset(
img_path=parser.voc_train + 'JPEGImages/',
xml_path=parser.voc_train + 'Annotations/',
class_list=class_list,
transform=transforms.Compose(
[Normalizer(), Augmenter(),
ResizerMultiScale()]))
if parser.voc_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = XML_VOCDataset(
img_path=parser.voc_val + 'JPEGImages/',
xml_path=parser.voc_val + 'Annotations/',
class_list=class_list,
transform=transforms.Compose([Normalizer(),
Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=1,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=2,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=2,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(),
pretrained=True)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=15,
verbose=True,
mode="max")
#scheduler = optim.lr_scheduler.StepLR(optimizer,8)
loss_hist = collections.deque(maxlen=1024)
retinanet.train()
retinanet.module.freeze_bn()
if not os.path.exists("./logs"):
os.mkdir("./logs")
log_file = open("./logs/log.txt", "w")
print('Num training images: {}'.format(len(dataset_train)))
best_map = 0
print("Training models...")
for epoch_num in range(parser.epochs):
#scheduler.step(epoch_num)
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
#print('iter num is: ', iter_num)
try:
#print(csv_eval.evaluate(dataset_val[:20], retinanet)[0])
#print(type(csv_eval.evaluate(dataset_val, retinanet)))
#print('iter num is: ', iter_num % 10 == 0)
optimizer.zero_grad()
classification_loss, regression_loss = retinanet(
[data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
#print(loss)
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
if iter_num % 50 == 0:
print(
'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'
.format(epoch_num, iter_num,
float(classification_loss),
float(regression_loss), np.mean(loss_hist)))
log_file.write(
'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f} \n'
.format(epoch_num, iter_num,
float(classification_loss),
float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, retinanet)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet)
elif parser.dataset == 'voc' and parser.voc_val is not None:
print('Evaluating dataset')
mAP = voc_eval.evaluate(dataset_val, retinanet)
try:
is_best_map = mAP[0][0] > best_map
best_map = max(mAP[0][0], best_map)
except:
pass
if is_best_map:
print("Get better map: ", best_map)
torch.save(retinanet.module,
'./logs/{}_scale15_{}.pt'.format(epoch_num, best_map))
shutil.copyfile(
'./logs/{}_scale15_{}.pt'.format(epoch_num, best_map),
"./best_models/model.pt")
else:
print("Current map: ", best_map)
scheduler.step(best_map)
retinanet.eval()
torch.save(retinanet, './logs/model_final.pt')
def train(csv_train=None, csv_classes=None, csv_val=None, epochs=12, depth=50, batch_size=2):
dataset = "csv"
# Create the data loaders
if dataset == 'csv':
if csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if csv_classes is None:
raise ValueError('Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(train_file=csv_train, class_list=csv_classes, transform=transforms.Compose([RandomHorizontalFlip(0.3),RandomRotation(6),Gamma_Correction(0.2), Image_Noise(0.2), Blur(0.2) , Normalizer(), Augmenter(), Resizer()]))
if csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=csv_val, class_list=csv_classes, transform=transforms.Compose([Normalizer(), Resizer()]))
else:
raise ValueError('Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train, batch_size=batch_size, drop_last=False)
dataloader_train = DataLoader(dataset_train, num_workers=3, collate_fn=collater, batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
dataloader_val = DataLoader(dataset_val, num_workers=3, collate_fn=collater, batch_sampler=sampler_val)
# Create the model
if depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(), pretrained=True)
elif depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(), pretrained=True)
elif depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(), pretrained=True)
elif depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(), pretrained=True)
elif depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(), pretrained=True)
else:
raise ValueError('Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
# Change
total_loss_data = []
class_loss_data = []
reg_loss_data = []
# Change
for epoch_num in range(epochs):
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
# Change
epoch_reg_loss = []
epoch_class_loss = []
# Change
for iter_num, data in enumerate(dataloader_train):
try:
optimizer.zero_grad()
classification_loss, regression_loss = retinanet([data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
# Change
epoch_reg_loss.append(float(regression_loss))
epoch_class_loss.append(float(classification_loss))
# Change
print('Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
if dataset == 'csv' and csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet)
# Change
total_loss_data.append(np.mean(epoch_loss))
class_loss_data.append(np.mean(epoch_class_loss))
reg_loss_data.append(np.mean(epoch_reg_loss))
print("Epoch loss", total_loss_data)
print("Epoch loss - classification", class_loss_data)
print("Epoch loss - Regression", reg_loss_data)
# Change
scheduler.step(np.mean(epoch_loss))
torch.save(retinanet.module, '{}_retinanet_{}.pt'.format(dataset, epoch_num))
retinanet.eval()
torch.save(retinanet, 'model_final.pt'.format(epoch_num))
# Change
import matplotlib.pyplot as plt
plt.plot(total_loss_data, label='Total loss')
plt.plot(class_loss_data, label='Classification loss')
plt.plot(reg_loss_data, label='Regression loss')
plt.ylabel("Loss")
plt.xlabel("Epoch")
plt.title("Epoch losses")
plt.legend()
plt.show()
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument(
'--train-file',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--classes-file',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--val-file',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=100)
parser.add_argument('--title', type=str, default='')
parser.add_argument("--resume_model", type=str, default="")
parser.add_argument("--resume_epoch", type=int, default=0)
parser.add_argument("--reinit-classifier",
action="store_true",
default=False)
parser.add_argument("--lr", type=float, default=.00001)
parser.add_argument("--all-box-regression",
action="store_true",
default=False)
parser.add_argument("--batch-size", type=int, default=16)
parser = parser.parse_args(args)
log_dir = "./runs/" + parser.title
writer = SummaryWriter(log_dir)
#pdb.set_trace()
with open(log_dir + '/config.csv', 'w') as f:
for item in vars(parser):
print(item + ',' + str(getattr(parser, item)))
f.write(item + ',' + str(getattr(parser, item)) + '\n')
if not os.path.isdir(log_dir + "/checkpoints"):
os.makedirs(log_dir + "/checkpoints")
if not os.path.isdir(log_dir + '/map_files'):
os.makedirs(log_dir + '/map_files')
dataset_train = CSVDataset(train_file=parser.train_file,
class_list=parser.classes_file,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
if parser.val_file is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.val_file,
class_list=parser.classes_file,
transform=transforms.Compose(
[Normalizer(), Resizer()]))
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=parser.batch_size,
drop_last=True)
dataloader_train = DataLoader(dataset_train,
num_workers=8,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=parser.batch_size,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=8,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(),
pretrained=True)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
if parser.resume_model:
x = torch.load(parser.resume_model)
if parser.reinit_classifier:
dummy = nn.Conv2d(256,
9 * dataset_train.num_classes(),
kernel_size=3,
padding=1)
x['classificationModel.output.weight'] = dummy.weight.clone()
x['classificationModel.output.bias'] = dummy.bias.clone()
prior = 0.01
x['classificationModel.output.weight'].data.fill_(0)
x['classificationModel.output.bias'].data.fill_(-math.log(
(1.0 - prior) / prior))
retinanet.load_state_dict(x)
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
#torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=parser.lr)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
# x = torch.load('./csv_retinanet_20.pth')
# retinanet.module.load_state_dict(x)
print('Num training images: {}'.format(len(dataset_train)))
for epoch_num in range(parser.resume_epoch, parser.epochs):
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
i = 0
avg_class_loss = 0.0
avg_reg_loss = 0.0
for iter_num, data in enumerate(dataloader_train):
i += 1
try:
optimizer.zero_grad()
#pdb.set_trace()
shape = data['img'].shape[2] * data['img'].shape[3]
writer.add_scalar("train/image_shape", shape,
epoch_num * (len(dataloader_train)) + i)
classification_loss, regression_loss = retinanet(
[data['img'].cuda().float(), data['annot'].cuda().float()])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
avg_class_loss += classification_loss
avg_reg_loss += regression_loss
if i % 100 == 0:
writer.add_scalar("train/classification_loss",
avg_class_loss / 100,
epoch_num * (len(dataloader_train)) + i)
writer.add_scalar("train/regression_loss",
avg_reg_loss / 100,
epoch_num * (len(dataloader_train)) + i)
avg_class_loss = 0.0
avg_reg_loss = 0.0
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
print(
'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'
.format(epoch_num, iter_num, float(classification_loss),
float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
if epoch_num % 2 == 0:
print('Evaluating dataset')
retinanet.eval()
mAP, AP_string = csv_eval.evaluate(dataset_val,
retinanet.module,
score_threshold=0.1)
with open(
log_dir + '/map_files/retinanet_{}.txt'.format(epoch_num),
'w') as f:
f.write(AP_string)
total = 0.0
all = 0.0
total_unweighted = 0.0
for c in mAP:
total += mAP[c][0] * mAP[c][1]
total_unweighted += mAP[c][0]
all += mAP[c][1]
writer.add_scalar("val/mAP", total / all, epoch_num)
writer.add_scalar("val/mAP_unweighted",
total_unweighted / len(mAP), epoch_num)
scheduler.step(np.mean(epoch_loss))
torch.save(retinanet.module.state_dict(),
log_dir + '/checkpoints/retinanet_{}.pth'.format(epoch_num))
retinanet.eval()
torch.save(retinanet.module.state_dict(),
log_dir + '/checkpoints/model_final.pth'.format(epoch_num))
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=20)
parser.add_argument('--resume',
'-r',
action='store_true',
help='resume from checkpoint')
parser.add_argument('--batch_size',
type=int,
default=4,
help='set the batch size')
parser.add_argument('--learning_rate',
'-lr',
type=float,
default=1e-3,
help='set the learning rate')
parser = parser.parse_args(args)
batch_size = parser.batch_size
learning_rate = parser.learning_rate
start_epoch = 0
device = 'cuda' if torch.cuda.is_available() else 'cpu'
best_acc = 0 # best test accuracy
# mean & std for ROI extraction
#mean = (0.1146, 0.1147, 0.1148)
#std = (0.1089, 0.1090, 0.1090)
# mean & std for QRCode extraction
mean = (0.2405, 0.2416, 0.2427)
std = (0.2194, 0.2208, 0.2223)
# Create the data loaders
if parser.dataset == 'coco':
if parser.coco_path is None:
raise ValueError('Must provide --coco_path when training on COCO,')
dataset_train = CocoDataset(parser.coco_path,
set_name='train2017',
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
dataset_val = CocoDataset(parser.coco_path,
set_name='val2017',
transform=transforms.Compose(
[Normalizer(), Resizer()]))
elif parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(
train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose([
Normalizer(mean=mean, std=std),
Augmenter(),
#YFlipAugmenter(),
#CropAugmenter(),
#Rot180Augmenter(),
Resizer()
]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose([
Normalizer(mean=mean, std=std),
Resizer()
]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=batch_size,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=4,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=4,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(),
pretrained=True)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
if use_gpu:
retinanet = retinanet.to(device)
retinanet = torch.nn.DataParallel(retinanet)
if parser.resume:
# Load checkpoint
print("==> Resuming from checkpoint")
checkpoint = torch.load('./checkpoint/ckpt.pth')
retinanet.load_state_dict(checkpoint['net'])
best_acc = checkpoint['acc']
start_epoch = checkpoint['epoch']
print('resume training from epoch:', start_epoch, " with accuracy:",
best_acc)
retinanet.training = True
#optimizer = optim.Adam(retinanet.parameters(), lr=1e-3)
optimizer = optim.SGD(retinanet.parameters(),
lr=learning_rate,
momentum=0.9)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
for epoch_num in range(start_epoch, start_epoch + parser.epochs):
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
try:
optimizer.zero_grad()
classification_loss, regression_loss = retinanet(
[data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
print(
'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'
.format(epoch_num, iter_num, float(classification_loss),
float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
writer.add_scalar('train_loss', np.mean(epoch_loss), epoch_num)
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, retinanet)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val,
retinanet,
iou_threshold=0.7,
max_detections=5,
score_threshold=0.2,
epoch=epoch_num)
print('mapROI:', mAP[0])
AP, num_annotations = mAP[0]
acc = 100. * AP
if acc > best_acc:
print('Saving... acc:', acc)
state = {
'net': retinanet.state_dict(),
'acc': acc,
'epoch': epoch_num,
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state, './checkpoint/ckpt.pth')
torch.save(retinanet, './checkpoint/best.pth')
best_acc = acc
writer.add_scalar('test_acc', acc, epoch_num)
scheduler.step(np.mean(epoch_loss))
#torch.save(retinanet.module, osp.join('checkpoints','{}_retinanet_{}.pt'.format(parser.dataset, epoch_num)))
retinanet.eval()
torch.save(retinanet, 'model_final.pt'.format(epoch_num))
writer.close()
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=100)
parser.add_argument('--title', type=str, default='')
parser.add_argument("--resume_model", type=str, default="")
parser.add_argument("--resume_epoch", type=int, default=0)
parser = parser.parse_args(args)
title = parser.resume_model.split('.')[0]
log_dir = "./runs/" + title
writer = SummaryWriter(log_dir)
if not os.path.isdir(log_dir + "/checkpoints"):
os.makedirs(log_dir + "/checkpoints")
if not os.path.isdir(log_dir + '/map_files'):
os.makedirs(log_dir + '/map_files')
if parser.dataset == 'csv':
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=0,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_val.num_classes(),
pretrained=True)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_val.num_classes(),
pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_val.num_classes(),
pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_val.num_classes(),
pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_val.num_classes(),
pretrained=True)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
if parser.resume_model:
retinanet.load_state_dict(torch.load(parser.resume_model))
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
theshes = [.05, 0.1, 0.2, 0.3]
i = 0
for thresh in theshes:
i = i + 1
retinanet.eval()
print('Evaluating dataset')
mAP, AP_string = csv_eval.evaluate(dataset_val,
retinanet,
score_threshold=thresh)
with open(
log_dir + '/map_files/{}_retinanet_{}.txt'.format(
parser.dataset, thresh), 'w') as f:
f.write(AP_string)
total = 0.0
all = 0.0
total_unweighted = 0.0
for c in mAP:
total += mAP[c][0] * mAP[c][1]
total_unweighted += mAP[c][0]
all += mAP[c][1]
writer.add_scalar("thresh_finder/mAP", total / all, i)
writer.add_scalar("thresh_finder/mAP_unweighted",
total_unweighted / len(mAP), i)
def main(args=None):
parser = argparse.ArgumentParser(
description='Training a RetinaNet network.')
parser.add_argument('--csv_train',
help='Path to file containing training annotations')
parser.add_argument('--csv_classes',
help='Path to file containing class list')
parser.add_argument('--csv_val',
help='Path to file containing validation \
annotations')
parser.add_argument("--depth",
help='Resnet depth, must be one of \
18, 34, 50,101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs to run',
type=int,
default=100)
parser.add_argument('--batch_size',
help='Number of training sample per batch',
type=int,
default=16)
parser.add_argument('--score_thresh',
help='score threshold to discard \
background/reduce nms processing time',
default=0.05)
parser.add_argument("--iou_nms1",
help="iou for nms used during validation and \
inference",
type=float,
default=0.3)
parser.add_argument('--lr', help='learning rate', type=float, default=6e-4)
parser.add_argument('--pretrained', type=bool, default=False)
parser.add_argument('--logfile')
args = parser.parse_args(args)
outputdir = os.path.dirname(args.logfile)
if not os.path.isdir(outputdir): os.makedirs(outputdir)
# Create the data loaders
if args.csv_train is None:
raise ValueError('Must provide --csv_train when training on CSV,')
if args.csv_classes is None:
raise ValueError('Must provide --csv_classes when training on CSV,')
dataset_train = CSVDataset(train_file=args.csv_train,
class_list=args.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
if args.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=args.csv_val,
class_list=args.csv_classes,
transform=transforms.Compose(
[Normalizer(), Resizer()]))
dataloader_train = DataLoader(dataset_train,
batch_size=args.batch_size,
num_workers=3,
collate_fn=collater,
shuffle=True)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if args.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(),
pretrained=args.pretrained)
elif args.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(),
pretrained=args.pretrained)
elif args.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(),
pretrained=args.pretrained)
elif args.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(),
pretrained=args.pretrained)
elif args.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(),
pretrained=args.pretrained)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
retinanet.score_thresh = args.score_thresh
retinanet.iou_nms1 = args.iou_nms1
optimizer = optim.Adam(retinanet.parameters(), lr=args.lr)
# # LR Finder
# lr_finder = LRFinder(retinanet, optimizer, losses.FocalLossQ, device="cuda")
# lr_finder.range_test(dataloader_train, end_lr=10, num_iter=1260, diverge_th=10)
# Ir_finder.plot(skip_start=0, skip_end=3, show_lr=3e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
loss_hist = collections.deque(maxlen=500)
print("Num training images: {}".format(len(dataset_train)))
for epoch_num in range(args.epochs):
retinanet.train()
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
optimizer.zero_grad()
classification_loss, regression_loss = retinanet(
[data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
print('Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | '
'Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(
epoch_num, iter_num, float(classification_loss),
float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
if args.csv_val is not None:
mAP = csv_eval.evaluate(dataset_val, retinanet)
with open(args.logfile, mode='a') as f:
f.write("mAP:\n")
aps = []
for i, label_name in enumerate(dataset_val.classes):
f.write('{}: {}| Count: {}\n'.format(
label_name, mAP[i][0], mAP[i][1]))
aps.append(mAP[i][0])
f.write('mAP: {}\n'.format(np.mean(aps)))
scheduler.step(np.mean(epoch_loss))
torch.save(retinanet.module,
'{}/retinanet_{}.pt'.format(outputdir, epoch_num))
torch.save(retinanet.module.state_dict(),
'{}/statedict_{}.pt'.format(outputdir, epoch_num))
retinanet.eval()
def main(args=None):
"""
In current implementation, if test csv is provided, we use that as validation set and combine the val and train csv's
as the csv for training.
If train_all_labeled_data flag is use, then we combine all 3 (if test is provided) for training and use a prespecified learning rate step schedule.
"""
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)',
default=None)
parser.add_argument(
'--csv_test',
help=
'Path to file containing test annotations (optional, if provided, train & val will be combined for training and test will be used for evaluation)',
default=None)
parser.add_argument('--lr', type=float, default=2e-5)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=101)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=25)
parser.add_argument('--model_output_dir', type=str, default='models')
parser.add_argument(
'--train_all_labeled_data',
help=
'Combine train, val, and test into 1 training set. Will use prespecified learning rate scheduler steps',
action='store_true')
parser.add_argument('--resnet-backbone-normalization',
choices=['batch_norm', 'group_norm'],
type=str,
default='batch_norm')
parser = parser.parse_args(args)
print('Learning Rate: {}'.format(parser.lr))
print("Normalization: ", parser.resnet_backbone_normalization)
# Create folder - will raise error if folder exists
assert (os.path.exists(parser.model_output_dir) == False)
os.mkdir(parser.model_output_dir)
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training,')
if parser.csv_classes is None:
raise ValueError('Must provide --csv_classes when training,')
if not parser.csv_val and parser.csv_test:
raise ValueError(
"Cannot specify test set without specifying validation set")
if parser.train_all_labeled_data:
csv_paths = [parser.csv_train, parser.csv_val, parser.csv_test]
train_csv = []
for path in csv_paths:
if isinstance(path, str):
train_csv.append(path)
val_csv = None
else:
if parser.csv_train and parser.csv_val and parser.csv_test:
train_csv = [parser.csv_train, parser.csv_val
] # Combine train and val sets for training
val_csv = parser.csv_test
else:
train_csv = parser.csv_train
val_csv = parser.csv_val
print('loading train data')
print(train_csv)
dataset_train = CSVDataset(train_file=train_csv,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
print(dataset_train.__len__())
if val_csv is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=val_csv,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(), Resizer()]))
print('putting data into loader')
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=2,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
print('creating model')
if parser.depth == 18:
retinanet = model.resnet18(
num_classes=dataset_train.num_classes(),
pretrained=True,
normalization=parser.resnet_backbone_normalization)
elif parser.depth == 34:
retinanet = model.resnet34(
num_classes=dataset_train.num_classes(),
pretrained=True,
normalization=parser.resnet_backbone_normalization)
elif parser.depth == 50:
retinanet = model.resnet50(
num_classes=dataset_train.num_classes(),
pretrained=True,
normalization=parser.resnet_backbone_normalization)
elif parser.depth == 101:
retinanet = model.resnet101(
num_classes=dataset_train.num_classes(),
pretrained=True,
normalization=parser.resnet_backbone_normalization)
elif parser.depth == 152:
retinanet = model.resnet152(
num_classes=dataset_train.num_classes(),
pretrained=True,
normalization=parser.resnet_backbone_normalization)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=parser.lr)
lr_factor = 0.3
if not parser.train_all_labeled_data:
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
factor=lr_factor,
verbose=True)
else:
# these milestones are for when using the lung masks - not for unmasked lung data
scheduler = optim.lr_scheduler.MultiStepLR(
optimizer, milestones=[12, 16, 20,
24], gamma=lr_factor) # masked training
#scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[14, 18, 22, 26], gamma=lr_factor)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
#initialize tensorboard
writer = SummaryWriter(comment=parser.model_output_dir)
# Augmentation
seq = iaa.Sequential([
iaa.Fliplr(0.5),
iaa.Flipud(0.5),
iaa.Affine(scale={
"x": (1.0, 1.2),
"y": (1.0, 1.2)
},
rotate=(-20, 20),
shear=(-4, 4))
],
random_order=True)
def augment(data, seq):
for n, img in enumerate(data['img']):
# imgaug needs dim in format (H, W, C)
image = data['img'][n].permute(1, 2, 0).numpy()
bbs_array = []
for ann in data['annot'][n]:
x1, y1, x2, y2, _ = ann
bbs_array.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2))
bbs = BoundingBoxesOnImage(bbs_array, shape=image.shape)
image_aug, bbs_aug = seq(image=image, bounding_boxes=bbs)
# save augmented image and chage dims to (C, H, W)
data['img'][n] = torch.tensor(image_aug.copy()).permute(2, 0, 1)
# save augmented annotations
for i, bbox in enumerate(bbs_aug.bounding_boxes):
x1, y1, x2, y2 = bbox.x1, bbox.y1, bbox.x2, bbox.y2
obj_class = data['annot'][n][i][-1]
data['annot'][n][i] = torch.tensor([x1, y1, x2, y2, obj_class])
return data
print('Num training images: {}'.format(len(dataset_train)))
dir_training_images = os.path.join(os.getcwd(), writer.log_dir,
'training_images')
os.mkdir(dir_training_images)
best_validation_loss = None
best_validation_map = None
for epoch_num in range(parser.epochs):
writer.add_scalar('Train/LR', optimizer.param_groups[0]['lr'],
epoch_num)
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
try:
optimizer.zero_grad()
data = augment(data, seq)
# save a few training images to see what augmentation looks like
if iter_num % 100 == 0 and epoch_num == 0:
x1, y1, x2, y2, _ = data['annot'][0][0]
fig, ax = plt.subplots(1)
ax.imshow(data['img'][0][1])
rect = patches.Rectangle((x1, y1),
x2 - x1,
y2 - y1,
linewidth=1,
edgecolor='r',
facecolor='none',
alpha=1)
ax.add_patch(rect)
fig.savefig(
os.path.join(dir_training_images,
'{}.png'.format(iter_num)))
plt.close()
classification_loss, regression_loss = retinanet(
[data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
if parser.resnet_backbone_normalization == 'batch_norm':
torch.nn.utils.clip_grad_norm_(
parameters=retinanet.parameters(), max_norm=0.1)
else:
torch.nn.utils.clip_grad_norm_(
parameters=retinanet.parameters(), max_norm=0.01
) # Decrease norm to reduce risk of exploding gradients
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
print(
'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'
.format(epoch_num, iter_num, float(classification_loss),
float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
writer.add_scalar('Train/Loss', np.mean(epoch_loss), epoch_num)
if not parser.train_all_labeled_data:
print('Evaluating Validation Loss...')
with torch.no_grad():
retinanet.train()
val_losses, val_class_losses, val_reg_losses = [], [], []
for val_iter_num, val_data in enumerate(dataloader_val):
try:
val_classification_loss, val_regression_loss = retinanet(
[
val_data['img'].cuda().float(),
val_data['annot']
])
val_losses.append(
float(val_classification_loss) +
float(val_regression_loss))
val_class_losses.append(float(val_classification_loss))
val_reg_losses.append(float(val_regression_loss))
del val_classification_loss, val_regression_loss
except Exception as e:
print(e)
continue
print(
'VALIDATION Epoch: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Total loss: {:1.5f}'
.format(epoch_num, np.mean(val_class_losses),
np.mean(val_reg_losses), np.mean(val_losses)))
# Save model with best validation loss
if best_validation_loss is None:
best_validation_loss = np.mean(val_losses)
if best_validation_loss >= np.mean(val_losses):
best_validation_loss = np.mean(val_losses)
torch.save(
retinanet.module,
parser.model_output_dir + '/best_result_valloss.pt')
writer.add_scalar('Validation/Loss', np.mean(val_losses),
epoch_num)
# Calculate Validation mAP
print('Evaluating validation mAP')
mAP = csv_eval.evaluate(dataset_val, retinanet)
print("Validation mAP: " + str(mAP[0][0]))
if best_validation_map is None:
best_validation_map = mAP[0][0]
elif best_validation_map < mAP[0][0]:
best_validation_map = mAP[0][0]
torch.save(
retinanet.module,
parser.model_output_dir + '/best_result_valmAP.pt')
writer.add_scalar('Validation/mAP', mAP[0][0], epoch_num)
if not parser.train_all_labeled_data:
scheduler.step(np.mean(val_losses))
else:
scheduler.step()
torch.save(
retinanet.module,
parser.model_output_dir + '/retinanet_{}.pt'.format(epoch_num))
retinanet.eval()
torch.save(retinanet, parser.model_output_dir + '/model_final.pt')
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--efficientdet',
help='Use EfficientDet.',
action="store_true")
parser.add_argument('--scaling-compound',
help='EfficientDet scaling compound phi.',
type=int,
default=0)
parser.add_argument('--batch-size', help='Batchsize.', type=int, default=6)
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument('--print-model-complexity',
help='Print model complexity.',
action="store_true")
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=None)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=100)
parser = parser.parse_args(args)
img_size = parser.scaling_compound * 128 + 512
# Create the data loaders
if parser.dataset == 'coco':
if parser.coco_path is None:
raise ValueError('Must provide --coco_path when training on COCO,')
dataset_train = CocoDataset(parser.coco_path,
set_name='train2017',
transform=transforms.Compose([
Normalizer(),
Augmenter(),
Resizer(img_size=img_size)
]))
dataset_val = CocoDataset(parser.coco_path,
set_name='val2017',
transform=transforms.Compose([
Normalizer(),
Resizer(img_size=img_size)
]))
elif parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose([
Normalizer(),
Augmenter(),
Resizer(img_size=img_size)
]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose([
Normalizer(),
Resizer(img_size=img_size)
]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=parser.batch_size,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
model = retinanet.resnet18(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 34:
model = retinanet.resnet34(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 50:
model = retinanet.resnet50(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 101:
model = retinanet.resnet101(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 152:
model = retinanet.resnet152(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.efficientdet:
model = efficientdet.efficientdet(
num_classes=dataset_train.num_classes(),
pretrained=True,
phi=parser.scaling_compound)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152, or specify '
)
use_gpu = True
if use_gpu:
model = model.cuda()
model = torch.nn.DataParallel(model).cuda()
if parser.print_model_complexity:
flops, params = get_model_complexity_info(model,
(3, img_size, img_size),
as_strings=True,
print_per_layer_stat=True)
print('{:<30} {:<8}'.format('Computational complexity: ', flops))
print('{:<30} {:<8}'.format('Number of parameters: ', params))
model.training = True
optimizer = optim.SGD(model.parameters(), lr=4e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
loss_hist = collections.deque(maxlen=500)
model.train()
model.module.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
for epoch_num in range(parser.epochs):
model.train()
model.module.freeze_bn()
freeze_layer(model.module.efficientnet)
epoch_loss = []
pbar = tqdm(enumerate(dataloader_train), total=len(dataloader_train))
for iter_num, data in pbar:
optimizer.zero_grad()
classification_loss, regression_loss = model(
[data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
mem = torch.cuda.memory_cached() / 1E9 if torch.cuda.is_available(
) else 0
pbar.set_description(
f'{mem:.3g}G | {float(classification_loss):1.5f} | {float(regression_loss):1.5f} | {np.mean(loss_hist):1.5f}'
)
#print('Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, model)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, model)
scheduler.step(np.mean(epoch_loss))
torch.save(model.module_model_,
'{}_model_{}.pt'.format(parser.dataset, epoch_num))
model.eval()
torch.save(model, 'model_final.pt'.format(epoch_num))
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=100)
parser = parser.parse_args(args)
# Create the data loaders
if parser.dataset == 'coco':
if parser.coco_path is None:
raise ValueError('Must provide --coco_path when training on COCO,')
# dataset_train = CocoDataset(parser.coco_path, set_name='trainval35k', transform=transforms.Compose([Normalizer(), Augmenter(), Resizer()]))
dataset_val = CocoDataset(parser.coco_path,
set_name='val5k',
transform=transforms.Compose(
[Normalizer(), Resizer()]))
elif parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
#sampler = AspectRatioBasedSampler(dataset_train, batch_size=16, drop_last=False)
#dataloader_train = DataLoader(dataset_train, num_workers=3, collate_fn=collater, batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_val.num_classes(),
pretrained=True)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_val.num_classes(),
pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_val.num_classes(),
pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_val.num_classes(),
pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_val.num_classes(),
pretrained=True)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
retinanet.load_state_dict(
torch.load("coco_resnet_50_map_0_335_state_dict.pt",
encoding='latin1'))
if use_gpu:
retinanet = retinanet.cuda()
# retinanet = torch.nn.DataParallel(retinanet).cuda()
#retinanet.training = True
#optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)
#scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True)
#loss_hist = collections.deque(maxlen=500)
retinanet.eval()
#retinanet.module.freeze_bn()
# print('Num training images: {}'.format(len(dataset_train)))
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, retinanet)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet)
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument(
'--dataset',
help='Dataset type, must be one of csv, coco or openimages')
parser.add_argument('--data_path', help='Path to COCO directory')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument('--resume', help='Checkpoint to load')
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=100)
parser.add_argument('--bs', help='Batch size', type=int, default=64)
parser.add_argument('--net', help='Network to use', default='fasterrcnn')
parser.add_argument('--log_interval',
help='Iterations before outputting stats',
type=int,
default=1)
parser.add_argument(
'--checkpoint_interval',
help='Iterations before saving an intermediate checkpoint',
type=int,
default=80)
parser.add_argument('--iterations',
type=int,
help='Iterations for every batch',
default=32)
parser = parser.parse_args(args)
# This becomes the minibatch size
parser.bs = parser.bs // parser.iterations
print('With {} iterations the effective batch size is {}'.format(
parser.iterations, parser.bs))
# Create the data loaders
if parser.dataset == 'coco':
raise NotImplementedError()
if parser.data_path is None:
raise ValueError('Must provide --data_path when training on COCO,')
dataset_train = CocoDataset(parser.data_path,
set_name='train2017',
transform=Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
dataset_val = CocoDataset(parser.data_path,
set_name='val2017',
transform=Compose([Normalizer(),
Resizer()]))
elif parser.dataset == 'csv':
raise NotImplementedError()
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=Compose(
[Normalizer(),
Resizer()]))
elif parser.dataset == 'openimages':
if parser.data_path is None:
raise ValueError(
'Must provide --data_path when training on OpenImages')
dataset_train = OidDataset(parser.data_path,
subset='train',
transform=Compose([
ToTensor(),
Augment(),
Resizer(min_side=600, max_side=1000)
]))
dataset_val = OidDataset(parser.data_path,
subset='validation',
transform=Compose([
ToTensor(),
Resizer(min_side=600, max_side=1000)
]))
elif parser.dataset == 'dummy':
# dummy dataset used only for debugging purposes
dataset_train = DummyDataset(
transform=Compose([ToTensor(), Resizer()]))
# dataset_val = DummyDataset(transform=Compose([ToTensor(), Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
sampler = BalancedSampler(dataset_train,
batch_size=parser.bs,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=8,
collate_fn=collate_fn,
batch_sampler=sampler)
# if dataset_val is not None:
# sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
# dataloader_val = DataLoader(dataset_val, num_workers=12, collate_fn=collate_fn, batch_sampler=sampler_val)
# Create the model
model = create_detection_model(dataset_train.num_classes(), parser)
# Create the experiment folder
experiment_fld = 'experiment_{}_{}_resnet{}_{}'.format(
parser.net, parser.dataset, parser.depth,
time.strftime("%Y%m%d%H%M%S", time.localtime()))
experiment_fld = os.path.join('outputs', experiment_fld)
if not os.path.exists(experiment_fld):
os.makedirs(experiment_fld)
logger = SummaryWriter(experiment_fld)
use_gpu = True
if use_gpu:
model = model.cuda()
model = torch.nn.DataParallel(model).cuda()
model.training = True
optimizer = optim.Adam(model.parameters(), lr=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
# Load checkpoint if needed
start_epoch = 0
if parser.resume:
print('Loading checkpoint {}'.format(parser.resume))
checkpoint = torch.load(parser.resume)
start_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['scheduler'])
print('Checkpoint loaded!')
loss_hist = collections.deque(maxlen=500)
model.train()
# model.module.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
for epoch_num in tqdm.trange(start_epoch, parser.epochs):
model.train()
# model.module.freeze_bn()
epoch_loss = []
log_losses_mean = {}
running_loss_sum = 0
data_progress = tqdm.tqdm(dataloader_train)
old_tensors_set = {}
optimizer.zero_grad()
for minibatch_idx, data in enumerate(data_progress):
images, targets = data
images = list(image.cuda().float() for image in images)
targets = [{k: v.cuda() for k, v in t.items()} for t in targets]
#images, targets = images.cuda(), targets.cuda()
loss_dict = model(images, targets)
#classification_loss = classification_loss.mean()
#regression_loss = regression_loss.mean()
#loss = classification_loss + regression_loss
loss = sum(loss for loss in loss_dict.values())
monitor_loss = loss.clone().detach()
loss /= parser.iterations
running_loss_sum += float(loss.item())
loss.backward()
if len(log_losses_mean) == 0:
log_losses_mean = clone_tensor_dict(loss_dict)
log_losses_mean['total_loss'] = float(monitor_loss.item())
else:
loss_dict['total_loss'] = monitor_loss
log_losses_mean = Counter(
clone_tensor_dict(loss_dict)) + Counter(log_losses_mean)
if (minibatch_idx + 1) % parser.iterations == 0:
data_progress.set_postfix(
dict(it=minibatch_idx // parser.iterations,
loss=running_loss_sum))
# all minibatches have been accumulated. Zero the grad
optimizer.step()
optimizer.zero_grad()
running_loss_sum = 0
# torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
# loss_hist.append(float(loss))
# epoch_loss.append(float(loss))
if (minibatch_idx + 1) % (parser.log_interval *
parser.iterations) == 0:
# compute the mean
log_losses_mean = {
k: (v / (parser.log_interval * parser.iterations))
for k, v in log_losses_mean.items()
}
logger.add_scalars(
"logs/losses", log_losses_mean,
epoch_num * len(dataloader_train) + minibatch_idx)
log_losses_mean = {}
# print('Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(loss_hist)))
if (minibatch_idx + 1) % (parser.checkpoint_interval *
parser.iterations) == 0:
# Save an intermediate checkpoint
save_checkpoint(
{
'model': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
'epoch': epoch_num
},
experiment_fld,
overwrite=True)
if (minibatch_idx + 1) % (1 * parser.iterations) == 0:
# flush cuda memory every tot iterations
torch.cuda.empty_cache()
'''for img in images:
del img
for tgt in targets:
for t in tgt.values():
del t
del loss'''
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, model)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, model)
# TODO: write evaluation code for openimages
scheduler.step(np.mean(epoch_loss))
save_checkpoint(
{
'model': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
'epoch': epoch_num,
},
experiment_fld,
overwrite=False)
model.eval()
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable
from torchvision import datasets, models, transforms
import torchvision
import model
from anchors import Anchors
import losses
from dataloader import CocoDataset, CSVDataset, collater, Resizer, AspectRatioBasedSampler, Augmenter, UnNormalizer, Normalizer
from torch.utils.data import Dataset, DataLoader
import coco_eval
import csv_eval
import warnings
warnings.filterwarnings("ignore")
assert torch.__version__.split('.')[1] == '4'
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
print('CUDA available: {}'.format(torch.cuda.is_available()))
dataset_val = CSVDataset(train_file="val.csv", class_list="classes.csv", transform=transforms.Compose([Normalizer(), Resizer()]))
retinanet = torch.load("./logs/csv_retinanet_139.pt").cuda()
retinanet.eval()
map = csv_eval.evaluate(dataset_val,retinanet)
print(map)
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument(
'--wider_train',
help='Path to file containing WIDER training annotations (see readme)')
parser.add_argument(
'--wider_val',
help=
'Path to file containing WIDER validation annotations (optional, see readme)'
)
parser.add_argument('--wider_train_prefix',
help='Prefix path to WIDER train images')
parser.add_argument('--wider_val_prefix',
help='Prefix path to WIDER validation images')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=50)
parser.add_argument('--batch_size',
help='Batch size (default 2)',
type=int,
default=2)
parser.add_argument('--model_name', help='Name of the model to save')
parser.add_argument('--parallel',
help='Run training with DataParallel',
dest='parallel',
default=False,
action='store_true')
parser.add_argument('--pretrained',
help='Pretrained model name in weight directory')
parser = parser.parse_args(args)
create_dirs()
# Create the data loaders
if parser.wider_train is None:
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Resizer(),
Augmenter(),
Normalizer()]))
else:
dataset_train = WIDERDataset(train_file=parser.wider_train,
img_prefix=parser.wider_train_prefix,
transform=transforms.Compose([
Resizer(),
Augmenter(),
Normalizer()
]))
if parser.wider_val is None:
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
print('Loading CSV validation dataset')
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Resizer(), Normalizer()]))
else:
print('Loading WIDER validation dataset')
dataset_val = WIDERDataset(train_file=parser.wider_val,
img_prefix=parser.wider_val_prefix,
transform=transforms.Compose(
[Resizer(), Normalizer()]))
print('Loading training dataset')
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=parser.batch_size,
drop_last=False)
if parser.parallel:
dataloader_train = DataLoader(dataset_train,
num_workers=16,
collate_fn=collater,
batch_sampler=sampler)
else:
dataloader_train = DataLoader(dataset_train,
collate_fn=collater,
batch_sampler=sampler)
# Create the model_pose_level_attention
if parser.depth == 18:
retinanet = resnet18(num_classes=dataset_train.num_classes())
elif parser.depth == 34:
retinanet = resnet34(num_classes=dataset_train.num_classes())
elif parser.depth == 50:
retinanet = resnet50(num_classes=dataset_train.num_classes())
elif parser.depth == 101:
retinanet = resnet101(num_classes=dataset_train.num_classes())
elif parser.depth == 152:
retinanet = resnet152(num_classes=dataset_train.num_classes())
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
if ckpt:
retinanet = torch.load('')
print('Loading checkpoint')
else:
print('Loading pretrained model')
retinanet_dict = retinanet.state_dict()
if parser.pretrained is None:
pretrained_dict = model_zoo.load_url(model_urls['resnet' +
str(parser.depth)])
else:
pretrained_dict = torch.load('./weight/' + parser.pretrained)
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in retinanet_dict
}
retinanet_dict.update(pretrained_dict)
retinanet.load_state_dict(retinanet_dict)
print('load pretrained backbone')
print(retinanet)
retinanet = torch.nn.DataParallel(retinanet, device_ids=[0])
if is_cuda:
retinanet.cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)
# optimizer = optim.SGD(retinanet.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
# scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
if parser.parallel:
retinanet.module.freeze_bn()
else:
retinanet.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
f_map = open('./mAP_txt/' + parser.model_name + '.txt', 'a')
writer = SummaryWriter(log_dir='./summary')
iters = 0
for epoch_num in range(0, parser.epochs):
retinanet.train()
if parser.parallel:
retinanet.module.freeze_bn()
else:
retinanet.freeze_bn()
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
iters += 1
optimizer.zero_grad()
img_data = data['img'].float()
annot_data = data['annot']
if is_cuda:
img_data = img_data.cuda()
annot_data = annot_data.cuda()
classification_loss, regression_loss, mask_loss = retinanet(
[img_data, annot_data])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
mask_loss = mask_loss.mean()
loss = classification_loss + regression_loss + mask_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
print(
'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | '
'mask_loss {:1.5f} | Running loss: {:1.5f}'.format(
epoch_num, iter_num, float(classification_loss),
float(regression_loss), float(mask_loss),
np.mean(loss_hist)))
writer.add_scalar('classification_loss',
float(classification_loss), iters)
writer.add_scalar('regression_loss', float(regression_loss), iters)
writer.add_scalar('loss', float(loss), iters)
del classification_loss
del regression_loss
del loss
if parser.wider_val is not None:
print('Evaluating dataset')
mAP = evaluate(dataset_val, retinanet, is_cuda=is_cuda)
f_map.write('mAP:{}, epoch:{}'.format(mAP[0][0], epoch_num))
f_map.write('\n')
scheduler.step(np.mean(epoch_loss))
if parser.parallel:
torch.save(
retinanet.module,
'./ckpt/' + parser.model_name + '_{}.pt'.format(epoch_num))
else:
torch.save(
retinanet,
'./ckpt/' + parser.model_name + '_{}.pt'.format(epoch_num))
retinanet.eval()
writer.export_scalars_to_json(
"./summary/' + parser.pretrained + 'all_scalars.json")
f_map.close()
writer.close()
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple testing script for RetinaNet network.')
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.',
default="csv")
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)',
default="binary_class.csv")
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--csv_box_annot',
help='Path to file containing predicted box annotations ')
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=18)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=500)
parser.add_argument('--model',
help='Path of .pt file with trained model',
default='esposallescsv_retinanet_0.pt')
parser.add_argument('--model_out',
help='Path of .pt file with trained model to save',
default='trained')
parser.add_argument('--score_threshold',
help='Score above which boxes are kept',
default=0.15)
parser.add_argument('--nms_threshold',
help='Score above which boxes are kept',
default=0.2)
parser.add_argument('--max_epochs_no_improvement',
help='Max epochs without improvement',
default=100)
parser.add_argument('--max_boxes',
help='Max boxes to be fed to recognition',
default=50)
parser.add_argument('--seg_level',
help='Line or word, to choose anchor aspect ratio',
default='line')
parser.add_argument(
'--htr_gt_box',
help='Train recognition branch with box gt (for debugging)',
default=False)
parser.add_argument(
'--binary_classifier',
help=
'Wether to use classification branch as binary or not, multiclass instead.',
default='False')
parser = parser.parse_args(args)
# Create the data loaders
if parser.dataset == 'csv':
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on COCO,')
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Resizer()]))
if parser.csv_box_annot is not None:
box_annot_data = CSVDataset(train_file=parser.csv_box_annot,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Resizer()]))
else:
box_annot_data = None
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=0,
collate_fn=collater,
batch_sampler=sampler_val)
if box_annot_data is not None:
sampler_val = AspectRatioBasedSampler(box_annot_data,
batch_size=1,
drop_last=False)
dataloader_box_annot = DataLoader(box_annot_data,
num_workers=0,
collate_fn=collater,
batch_sampler=sampler_val)
else:
dataloader_box_annot = dataloader_val
if not os.path.exists('trained_models'):
os.mkdir('trained_models')
# Create the model
alphabet = dataset_val.alphabet
if os.path.exists(parser.model):
retinanet = torch.load(parser.model)
else:
print("Choose an existing saved model path.")
sys.exit()
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
#retinanet = torch.load('../Documents/TRAINED_MODELS/pytorch-retinanet/esposallescsv_retinanet_99.pt')
#print "LOADED pretrained MODEL\n\n"
optimizer = optim.Adam(retinanet.parameters(), lr=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=4,
verbose=True)
loss_hist = collections.deque(maxlen=500)
ctc = CTCLoss()
retinanet.module.freeze_bn()
best_cer = 1000
epochs_no_improvement = 0
cers = []
retinanet.eval()
retinanet.module.epochs_only_det = 0
#retinanet.module.htr_gt_box = False
retinanet.training = False
if parser.score_threshold is not None:
retinanet.module.score_threshold = float(parser.score_threshold)
'''if parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
'''
mAP, binary_mAP, cer = csv_eval.evaluate(
dataset_val,
retinanet,
score_threshold=retinanet.module.score_threshold)
def train(img_dir,classes_csv,model_fname=None,resnet_depth=50,epochs=1000,steps=100,train_split=0.8,out_dir ='',out_prefix=''):
if not os.path.exists(out_dir):
os.makedirs(out_dir)
# Create the data loaders
# Get all image fnames in folder
img_list = []
if not isinstance(img_dir, list):
img_dir = [img_dir]
for dir in img_dir:
for file in os.listdir(dir):
if file.endswith(".png"):
img_list.append(dir + file)
randomised_list = random.sample(img_list, len(img_list))
num_train = int(0.8*len(img_list))
train_imgs, val_imgs = randomised_list[:num_train], randomised_list[num_train:]
dataset_train = CustomDataset(img_list=train_imgs, class_list=classes_csv, transform=transforms.Compose([Normalizer(), Augmenter(), Resizer()]))
dataset_val = CustomDataset(img_list=val_imgs, class_list=classes_csv,transform=transforms.Compose([Normalizer(), Resizer()]))
sampler = AspectRatioBasedSampler(dataset_train, batch_size=2, drop_last=False)
dataloader_train = DataLoader(dataset_train, num_workers=3, collate_fn=collater, batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
dataloader_val = DataLoader(dataset_val, num_workers=3, collate_fn=collater, batch_sampler=sampler_val)
# Create the model
if resnet_depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(), pretrained=True)
elif resnet_depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(), pretrained=True)
elif resnet_depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(), pretrained=True)
elif resnet_depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(), pretrained=True)
elif resnet_depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(), pretrained=True)
else:
raise ValueError('Unsupported model depth, must be one of 18, 34, 50, 101, 152')
# retinanet = torch.load(model_fname)
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
if model_fname is not None:
retinanet.load_state_dict(torch.load(model_fname))
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
start_time = time.clock()
print('Num training images: {}'.format(len(dataset_train)))
for epoch_num in range(epochs):
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
try:
optimizer.zero_grad()
classification_loss, regression_loss = retinanet([data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
# print('Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
print('Epoch: {} | Running loss: {:1.5f} | Elapsed Time: {}'.format(epoch_num, np.mean(loss_hist),(time.clock() - start_time)/60))
mAP = csv_eval.evaluate(dataset_val, retinanet)
scheduler.step(np.mean(epoch_loss))
if (epoch_num) % steps == 0:
torch.save(retinanet.module, '{}{}_model_{}.pt'.format(out_dir, out_prefix, epoch_num))
torch.save(retinanet.state_dict(), '{}{}_state_{}.pt'.format(out_dir, out_prefix, epoch_num))
torch.save(retinanet, out_dir + '{}model_final.pt'.format(out_prefix))
torch.save(retinanet.state_dict(), out_dir + '{}state_final_.pt'.format(out_prefix))
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument('--model', help='Path to model (.pt) file.')
parser.add_argument("--save_path", help="Path to save evaluated images")
parser = parser.parse_args(args)
if parser.dataset == 'coco':
dataset_val = CocoDataset(parser.coco_path,
set_name='val2017',
transform=transforms.Compose(
[Normalizer(), Resizer()]))
elif parser.dataset == 'csv':
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(), Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=1,
collate_fn=collater,
batch_sampler=sampler_val)
retinanet = torch.load(parser.model)
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
retinanet.eval()
mAP, carmAP, smokemAP = csv_eval.evaluate(dataset_val, retinanet)
unnormalize = UnNormalizer()
def draw_caption(image, box, caption):
b = np.array(box).astype(int)
cv2.putText(image, caption, (b[0], b[1] - 10), cv2.FONT_HERSHEY_PLAIN,
1, (0, 0, 0), 2)
cv2.putText(image, caption, (b[0], b[1] - 10), cv2.FONT_HERSHEY_PLAIN,
1, (255, 255, 255), 1)
times = []
for idx, data in enumerate(dataloader_val):
with torch.no_grad():
st = time.time()
scores, classification, transformed_anchors = retinanet(
data['img'].cuda().float())
times.append(time.time() - st)
idxs = np.where(scores > 0.5)
img = np.array(255 * unnormalize(data['img'][0, :, :, :])).copy()
img[img < 0] = 0
img[img > 255] = 255
img = np.transpose(img, (1, 2, 0))
img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_BGR2RGB)
for j in range(idxs[0].shape[0]):
bbox = transformed_anchors[idxs[0][j], :]
x1 = int(bbox[0])
y1 = int(bbox[1])
x2 = int(bbox[2])
y2 = int(bbox[3])
label_name = dataset_val.labels[int(
classification[idxs[0][j]])]
draw_caption(img, (x1, y1, x2, y2), label_name)
cv2.rectangle(img, (x1, y1), (x2, y2),
color=(0, 0, 255),
thickness=2)
cv2.imwrite(os.path.join(parser.save_path, f'img_{j}.png'), img)
cv2.waitKey(0)
print("Average Elapsed Time: {}".format(np.mean(times)))
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=50)
parser.add_argument('--model_name', help='name of the model to save')
parser.add_argument('--pretrained', help='pretrained model name')
parser = parser.parse_args(args)
# Create the data loaders
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Resizer(),
Augmenter(),
Normalizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Resizer(), Normalizer()]))
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=2,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=16,
collate_fn=collater,
batch_sampler=sampler)
#dataloader_train = DataLoader(dataset_train, num_workers=16, collate_fn=collater, batch_size=8, shuffle=True)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=2,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=16,
collate_fn=collater,
batch_sampler=sampler_val)
#dataloader_val = DataLoader(dataset_train, num_workers=16, collate_fn=collater, batch_size=8, shuffle=True)
# Create the model_pose_level_attention
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes())
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes())
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes())
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes())
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes())
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
if ckpt:
retinanet = torch.load('')
print('load ckpt')
else:
retinanet_dict = retinanet.state_dict()
pretrained_dict = torch.load('./weight/' + parser.pretrained)
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in retinanet_dict
}
retinanet_dict.update(pretrained_dict)
retinanet.load_state_dict(retinanet_dict)
print('load pretrained backbone')
print(retinanet)
retinanet = torch.nn.DataParallel(retinanet, device_ids=[0])
retinanet.cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)
#optimizer = optim.SGD(retinanet.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
#scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
f_map = open('./mAP_txt/' + parser.model_name + '.txt', 'a')
writer = SummaryWriter(log_dir='./summary')
iters = 0
for epoch_num in range(0, parser.epochs):
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
#scheduler.step()
for iter_num, data in enumerate(dataloader_train):
iters += 1
optimizer.zero_grad()
classification_loss_f, regression_loss_f, classification_loss_v, regression_loss_v = retinanet(
[
data['img'].cuda().float(), data['annot'], data['vbox'],
data['ignore']
])
classification_loss_f = classification_loss_f.mean()
regression_loss_f = regression_loss_f.mean()
classification_loss_v = classification_loss_v.mean()
regression_loss_v = regression_loss_v.mean()
loss = classification_loss_f + regression_loss_f + classification_loss_v + regression_loss_v
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
print(
'Epoch: {} | Iteration: {} | Classification loss_f: {:1.5f} | Regression loss_f: {:1.5f} | Classification loss_v {:1.5f} | Regression loss_v {:1.5f} | Running loss: {:1.5f}'
.format(epoch_num, iter_num, float(classification_loss_f),
float(regression_loss_f), float(classification_loss_v),
float(regression_loss_v), np.mean(loss_hist)))
writer.add_scalar('classification_loss_f', classification_loss_f,
iters)
writer.add_scalar('regression_loss_f', regression_loss_f, iters)
writer.add_scalar('classification_loss_v', classification_loss_v,
iters)
writer.add_scalar('regression_loss_v', regression_loss_v, iters)
writer.add_scalar('loss', loss, iters)
if parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet)
f_map.write('mAP:{}, epoch:{}'.format(mAP[0][0], epoch_num))
f_map.write('\n')
scheduler.step(np.mean(epoch_loss))
torch.save(retinanet.module,
'./ckpt/' + parser.model_name + '_{}.pt'.format(epoch_num))
retinanet.eval()
writer.export_scalars_to_json(
"./summary/' + parser.pretrained + 'all_scalars.json")
f_map.close()
writer.close()
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path', help='Path to COCO directory')
parser.add_argument(
'--csv_train',
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=50)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=100)
parser = parser.parse_args(args)
# Create the data loaders
if parser.dataset == 'coco':
if parser.coco_path is None:
raise ValueError('Must provide --coco_path when training on COCO,')
dataset_train = CocoDataset(parser.coco_path,
set_name='train2017',
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
dataset_val = CocoDataset(parser.coco_path,
set_name='val2017',
transform=transforms.Compose(
[Normalizer(), Resizer()]))
elif parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=2,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(),
pretrained=True)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
for epoch_num in range(parser.epochs):
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
try:
optimizer.zero_grad()
classification_loss, regression_loss = retinanet(
[data['img'].cuda().float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
print(
'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'
.format(epoch_num, iter_num, float(classification_loss),
float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
except Exception as e:
print(e)
continue
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, retinanet)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet)
scheduler.step(np.mean(epoch_loss))
torch.save(
retinanet.module,
'{}_retinanet_dilation_{}.pt'.format(parser.dataset, epoch_num))
retinanet.eval()
torch.save(retinanet, 'model_final_dilation.pt'.format(epoch_num))
def main(args=None):
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',
default="csv",
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--coco_path',
default="/home/mayank-s/PycharmProjects/Datasets/coco",
help='Path to COCO directory')
parser.add_argument(
'--csv_train',
default="berkely_ready_to_train_for_retinanet_pytorch.csv",
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
default="berkely_class.csv",
help='Path to file containing class list (see readme)')
parser.add_argument(
'--csv_val',
default=
"/home/teai/Desktop/mayank/pytorch-retinanet-master (4)/berkely_ready_to_train_validation.csv",
help=
'Path to file containing validation annotations (optional, see readme)'
)
# parser.add_argument('--csv_val',default="/home/teai/Desktop/mayank/pytorch-retinanet-master (4)/berkely_ready_to_train_validation (copy).csv", help='Path to file containing validation annotations (optional, see readme)')
parser.add_argument(
'--depth',
help='Resnet depth, must be one of 18, 34, 50, 101, 152',
type=int,
default=18)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=200)
parser = parser.parse_args(args)
# Create the data loaders
if parser.dataset == 'coco':
if parser.coco_path is None:
raise ValueError('Must provide --coco_path when training on COCO,')
dataset_train = CocoDataset(parser.coco_path,
set_name='train2014',
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
dataset_val = CocoDataset(parser.coco_path,
set_name='val2014',
transform=transforms.Compose(
[Normalizer(), Resizer()]))
elif parser.dataset == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes is None:
raise ValueError(
'Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
if parser.csv_val is None:
dataset_val = None
print('No validation annotations provided.')
else:
dataset_val = CSVDataset(train_file=parser.csv_val,
class_list=parser.csv_classes,
transform=transforms.Compose(
[Normalizer(),
Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=4,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=0,
collate_fn=collater,
batch_sampler=sampler)
if dataset_val is not None:
sampler_val = AspectRatioBasedSampler(dataset_val,
batch_size=1,
drop_last=False)
dataloader_val = DataLoader(dataset_val,
num_workers=3,
collate_fn=collater,
batch_sampler=sampler_val)
# Create the model
if parser.depth == 18:
retinanet = model.resnet18(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 34:
retinanet = model.resnet34(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 50:
retinanet = model.resnet50(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 101:
retinanet = model.resnet101(num_classes=dataset_train.num_classes(),
pretrained=True)
elif parser.depth == 152:
retinanet = model.resnet152(num_classes=dataset_train.num_classes(),
pretrained=True)
else:
raise ValueError(
'Unsupported model depth, must be one of 18, 34, 50, 101, 152')
use_gpu = True
# if use_gpu:
if torch.cuda.is_available():
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
# checkpoint='/home/teai/Desktop/mayank/pytorch-retinanet-master (4)/checkpoint/old_datasets_with_missing_images_annotations/retina_fpn_11'
checkpoint = '/home/teai/Desktop/mayank/pytorch-retinanet-master (4)/checkpoint/retina_fpn_univ_132'
# retinanet = torch.load("./checkpoint/retina_fpn_132")
retinanet = torch.load(checkpoint)
for epoch_num in range(parser.epochs):
# retinanet.train()retina_fpn_1
# retinanet.module.freeze_bn()
epoch_loss = []
'''for iter_num, data in enumerate(dataloader_train):
try:
optimizer.zero_grad()
if torch.cuda.is_available():
classification_loss, regression_loss = retinanet([data['img'].cuda().float(), data['annot']])
else:
classification_loss, regression_loss = retinanet([data['img'].float(), data['annot']])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
loss = classification_loss + regression_loss
if bool(loss == 0):
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
optimizer.step()
loss_hist.append(float(loss))
epoch_loss.append(float(loss))
print('Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(loss_hist)))
del classification_loss
del regression_loss
except Exception as e:
print(e)
print("fail")
continue
print("Saving model...")
name = "./checkpoint/retina_fpn_" + str(epoch_num)
torch.save(retinanet, name)'''
if parser.dataset == 'coco':
print('Evaluating dataset')
coco_eval.evaluate_coco(dataset_val, retinanet)
elif parser.dataset == 'csv' and parser.csv_val is not None:
print('Evaluating dataset')
mAP = csv_eval.evaluate(dataset_val, retinanet)
# print(mAP)
scheduler.step(np.mean(epoch_loss))
'''
