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 train(args):
train_csv = args.train_csv
test_csv = args.test_csv
labels_csv = args.labels_csv
model_type = args.model_type
epochs = int(args.epochs)
batch_size = int(args.batch_size)
dataset_train = CSVDataset(train_file=train_csv,
class_list=labels_csv,
transform=transforms.Compose(
[Normalizer(),
Augmenter(),
Resizer()]))
dataset_val = CSVDataset(train_file=test_csv,
class_list=labels_csv,
transform=transforms.Compose(
[Normalizer(), Resizer()]))
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)
retinanet = RetinaNet_efficientnet_b4(
num_classes=dataset_train.num_classes(), model_type=model_type)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
retinanet = retinanet.to(device)
retinanet = torch.nn.DataParallel(retinanet).to(device)
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(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, regression_loss = retinanet(
[data['img'].to(device).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
#mAP, MAP = evaluate(dataset_val, retinanet)
_, MAP = evaluate(dataset_val, retinanet)
scheduler.step(np.mean(epoch_loss))
torch.save(
retinanet.module,
'{}_retinanet_{}_map{}.pt'.format("EfficientNet" + model_type,
epoch_num, MAP))
retinanet.eval()
torch.save(retinanet, 'model_final.pt')
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()
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_experiment_{}.pt'.format(
parser.dataset, epoch_num))
name = '{}_retinanet_dilation_experiment_{}.pt'.format(
parser.dataset, epoch_num)
parser.resume = '/users/wenchi/ghwwc/pytorch-retinanet-master_new/name'
retinanet.eval()
torch.save(retinanet,
'model_final_dilation_experiment.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', 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('--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()]))
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=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=5, verbose=True,mode="max")
#scheduler = optim.lr_scheduler.StepLR(optimizer,8)
loss_hist = collections.deque(maxlen=500)
retinanet.train()
retinanet.module.freeze_bn()
if not os.path.exists("./logs"):
os.mkdir("./logs")
if not os.path.exists('best_models'):
os.makedirs('best_models')
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):
try:
#print(csv_eval.evaluate(dataset_val[:20], retinanet)[0])
#print(type(csv_eval.evaluate(dataset_val, retinanet)))
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))
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)
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 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(
'--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 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))
'''
def __init__(self):
self.dataset = 'csv'
self.csv_train = './annotations_train.csv'
self.csv_classes = './classes.csv'
self.depth = 50
self.epochs = 3
self.csv_val = './annotations_val.csv'
parser = arguments()
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()]))
sampler = AspectRatioBasedSampler(dataset_train, batch_size=2, drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=3,
collate_fn=collater,
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='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):
from dataloader import JinNanDataset, Augmenter, UnNormalizer, Normalizer,Resizer
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, models, transforms
import model
import torch
import argparse
parser = argparse.ArgumentParser(description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset',default='jingnan', help='Dataset type, must be one of csv or coco.')
parser.add_argument('--threshold',help='treshold')
parser.add_argument('--dataset_path', help='Path to file containing training and validation annotations (optional, see readme)')
parser.add_argument('--model_path',help=('the model path'))
parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152', type=int, default=50)
parser = parser.parse_args(args)
dataset_val=JinNanDataset(parser.dataset_path, set_name='val', transform=transforms.Compose([Normalizer(), Augmenter(), Resizer()]))
# 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')
retinanet=torch.load(parser.model_path)
use_gpu = True
if use_gpu:
retinanet = retinanet.cuda()
retinanet.eval()
print('Evaluating dataset')
evaluate_jinnan(dataset_val, retinanet)
val_batch_size = 4
num_workers = 3
lr=1e-5
patience=3
verbose=True
maxlen=500
threshold=0.05
model_name = 'retinanet'
epochs = 10
# Load train and validation dataset (for sake of example i have used same but use different dataset)
# Load train image folder and corresponding coco json file to train dataset
# Load validation image folder and corresponding json file to validation dataset
images_folder = '/content/data/val2017'
train_json_file = '/content/data/train_coco_dataset.json'
val_json_file = '/content/data/val_coco_dataset.json'
dataset_train = CocoDataset(images_folder, train_json_file,
transform=transforms.Compose([Normalizer(), Augmenter(), Resizer()]))
dataset_val = CocoDataset(images_folder, val_json_file,
transform=transforms.Compose([Normalizer(), Resizer()]))
num_classes = dataset_train.num_classes()
# run main function
main()
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 = torch.cuda.is_available()
if use_gpu:
efficientdet = efficientdet.cuda()
efficientdet = torch.nn.DataParallel(efficientdet).cuda()
efficientdet.training = True
optimizer = optim.Adam(efficientdet.parameters(), lr=1e-5)
# TODO: Add learning rate scheduler
# Calculate 5% of warm-up training
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
loss_hist = collections.deque(maxlen=500)
efficientdet.train()
efficientdet.module.freeze_bn()
print(
f"Number of parameters: {sum(p.numel() for p in efficientdet.parameters())}"
)
print(
f"Number of trainable parameters: {sum(p.numel() for p in efficientdet.parameters() if p.requires_grad)}"
)
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.')
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_general',
help='Path to file containing class list (see readme)')
parser.add_argument('--csv_features',
help='Path to dir containing features csv files')
parser.add_argument('--csv_colors',
help='Path to file containing color classes')
parser.add_argument('--csv_types',
help='Path to file containing type classes')
parser.add_argument(
'--csv_val',
help=
'Path to file containing validation annotations (optional, see readme)'
)
parser.add_argument(
'--image_dir',
help='Path to file containing images (optional, see readme)')
parser.add_argument('--pretrain_model', help='Path to model (.pt) file.')
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 == 'csv':
if parser.csv_train is None:
raise ValueError('Must provide --csv_train when training on COCO,')
if parser.csv_classes_general is None:
raise ValueError(
'Must provide --csv_classes_general when training on COCO,')
dataset_train = CSVDataset(train_file=parser.csv_train,
class_list=parser.csv_classes_general,
color_classes=parser.csv_colors,
type_classes=parser.csv_types,
feature_class_dir=parser.csv_features,
image_dir=parser.image_dir,
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_general,
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 = nn.DataParallel(retinanet)
# torch.cuda.set_device(0)
retinanet = retinanet.cuda()
retinanet = torch.nn.DataParallel(retinanet).cuda()
retinanet.training = True
if parser.pretrain_model is not None:
retinanet = torch.load(parser.pretrain_model)
print('load model: ' + str(parser.pretrain_model))
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_{}.pt'.format(parser.dataset, epoch_num))
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(
'--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 CTracker network.')
parser.add_argument('--dataset',
default='csv',
type=str,
help='Dataset type, must be one of csv or coco.')
parser.add_argument('--model_dir',
default='./ctracker/',
type=str,
help='Path to save the model.')
parser.add_argument(
'--root_path',
default='/dockerdata/home/changanwang/Dataset/Tracking/MOT17Det/',
type=str,
help='Path of the directory containing both label and images')
parser.add_argument(
'--csv_train',
default='train_annots.csv',
type=str,
help='Path to file containing training annotations (see readme)')
parser.add_argument('--csv_classes',
default='train_labels.csv',
type=str,
help='Path to file containing class list (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)
print(parser)
print(parser.model_dir)
if not os.path.exists(parser.model_dir):
os.makedirs(parser.model_dir)
# Create the data loaders
if parser.dataset == 'csv':
if (parser.csv_train is None) or (parser.csv_train == ''):
raise ValueError('Must provide --csv_train when training on COCO,')
if (parser.csv_classes is None) or (parser.csv_classes == ''):
raise ValueError(
'Must provide --csv_classes when training on COCO,')
dataset_train = CSVDataset(parser.root_path, train_file=os.path.join(parser.root_path, parser.csv_train), class_list=os.path.join(parser.root_path, parser.csv_classes), \
transform=transforms.Compose([RandomSampleCrop(), PhotometricDistort(), Augmenter(), Normalizer()]))#transforms.Compose([Normalizer(), Augmenter(), Resizer()]))
else:
raise ValueError(
'Dataset type not understood (must be csv or coco), exiting.')
# sampler = AspectRatioBasedSampler(dataset_train, batch_size=2, drop_last=False)
sampler = AspectRatioBasedSampler(dataset_train,
batch_size=8,
drop_last=False)
dataloader_train = DataLoader(dataset_train,
num_workers=32,
collate_fn=collater,
batch_sampler=sampler)
# 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)
optimizer = optim.Adam(retinanet.parameters(), lr=5e-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)))
total_iter = 0
for epoch_num in range(parser.epochs):
retinanet.train()
retinanet.module.freeze_bn()
epoch_loss = []
for iter_num, data in enumerate(dataloader_train):
try:
total_iter = total_iter + 1
optimizer.zero_grad()
(classification_loss, regression_loss), reid_loss = retinanet([
data['img'].cuda().float(), data['annot'],
data['img_next'].cuda().float(), data['annot_next']
])
classification_loss = classification_loss.mean()
regression_loss = regression_loss.mean()
reid_loss = reid_loss.mean()
# loss = classification_loss + regression_loss + track_classification_losses
loss = classification_loss + regression_loss + reid_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: {} | Iter: {} | Cls loss: {:1.5f} | Reid loss: {:1.5f} | Reg loss: {:1.5f} | Running loss: {:1.5f}'
.format(epoch_num, iter_num, float(classification_loss),
float(reid_loss), float(regression_loss),
np.mean(loss_hist)))
except Exception as e:
print(e)
continue
scheduler.step(np.mean(epoch_loss))
retinanet.eval()
torch.save(retinanet, os.path.join(parser.model_dir, 'model_final.pt'))
run_from_train(parser.model_dir, parser.root_path)
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('--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=152)
parser.add_argument('--epochs',
help='Number of epochs',
type=int,
default=100)
parser.add_argument("--save_path", help="save path", type=str)
parser = parser.parse_args(args)
## Save path
if not os.path.exists(parser.save_path):
os.makedirs(parser.save_path)
# 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 = False
if use_gpu:
retinanet = retinanet.cuda()
# Commenting data parallel
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()
#retinanet.freeze_bn()
print('Num training images: {}'.format(len(dataset_train)))
print("Epochs:", parser.epochs)
epoch_losses = []
car_map_values = []
smoke_map_values = []
print(parser.save_path)
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()
if iter_num == 1:
print(data['img'].size())
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, carmAP, smokemAP = csv_eval.evaluate(dataset_val, retinanet)
scheduler.step(np.mean(epoch_loss))
epoch_losses.append(np.mean(epoch_loss))
car_map_values.append(carmAP)
smoke_map_values.append(smokemAP)
torch.save(
retinanet.module,
os.path.join(
parser.save_path,
'{}_retinanet_{}.pt'.format(parser.dataset, epoch_num)))
retinanet.eval()
torch.save(retinanet, 'model_final.pt'.format(epoch_num))
plt.plot(list(range(parser.epochs)), epoch_losses)
plt.plot(list(range(parser.epochs)), car_map_values)
plt.plot(list(range(parser.epochs)), smoke_map_values)
plt.title("Training Epoch loss Vs mAP Cars")
plt.xlabel("epochs")
plt.ylabel("value")
plt.legend(["Loss", "Car mAP", "Smoke mAP"])
plt.savefig("train_loss_mAP_smoke.png")
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('--basemodel', help='Number of epochs', default='')
parser.add_argument('--save_path', help='Path to save output model.', default='./')
parser = parser.parse_args(args)
# 接受并创建输出路径
output_path = parser.save_path
if os.path.exists(output_path):
print('Output path exist: %s.' % output_path)
else:
print('Creating output path: %s' % output_path)
os.mkdir(output_path)
# 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=8, 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=8, collate_fn=collater, batch_sampler=sampler_val)
if len(parser.basemodel) > 0 and os.path.exists(parser.basemodel):
print('restore model from %s' % parser.basemodel)
retinanet = torch.load(parser.basemodel)
#retinanet = model.resnet18(num_classes=dataset_train.num_classes(), pretrained=True)
else:
# 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):
#continue
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):
print(data['img'].shape, data['annot'])
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)))
if iter_num % 1000 == 0 and iter_num >0:
print('saving tmp model.')
retinanet.eval()
torch.save(retinanet, os.path.join(output_path, 'model_tmp_%d.pt' % iter_num))
retinanet.train()
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, os.path.join(output_path, '{}_retinanet_{}.pt'.format(parser.dataset, epoch_num)))
retinanet.eval()
torch.save(retinanet, os.path.join(output_path, 'model_final.pt'.format(epoch_num)))