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