def test_voc():
# load net
num_classes = len(VOC_CLASSES) + 1 # +1 background
net = build_ssd('test', 300, num_classes) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
# load data
testset = VOCDetection(args.voc_root, [('2007', 'test')], None, VOCAnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(args.save_folder, net, args.cuda, testset,
BaseTransform(net.size, (104, 117, 123)),
thresh=args.visual_threshold)
# detections = [batch, number of classes, number of occurence, (score, x0, Y0, x1, y1)]
for i in range(detections.size(1)):
j = 0
while detections[0, i, j, 0] >= 0.6:
pt = (detections[0, i, j, 1:] * scale).numpy()
cv2.rectangle(frame, (int(pt[0]), int(pt[1])),
(int(pt[2]), int(pt[3])), (255, 0, 0), 2)
cv2.putText(frame, labelmap[i - 1], (int(pt[0]), int(pt[1])),
cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 2,
cv2.LINE_AA)
j += 1
return frame
# Creating the SSD neural network
net = build_ssd('test')
net.load_state_dict(
torch.load('ssd300_mAP_77.43_v2.pth',
map_location=lambda storage, loc: storage))
# Creating the transformation
transform = BaseTransform(net.size, (104 / 256.0, 117 / 256.0, 123 / 256.0))
# Doing some Object Detection on a video
reader = imageio.get_reader('funny_dog.mp4')
fps = reader.get_meta_data()['fps']
writer = imageio.get_writer('output.mp4', fps=fps)
for i, frame in enumerate(reader):
frame = detect(frame, net.eval(), transform)
writer.append_data(frame)
print(i)
torch.set_default_tensor_type('torch.cuda.FloatTensor')
if not args.cuda:
print(
"WARNING: It looks like you have a CUDA device, but aren't using \ CUDA. Run with --cuda for optimal eval speed."
)
torch.set_default_tensor_type('torch.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
#if torch.cuda.is_available():
# if args.cuda:
# torch.set_default_tensor_type('torch.cuda.FloatTensor')
# if not args.cuda:
# torch.set_default_tensor_type('torch.FloatTensor')
net = build_ssd('test', 300, 21) # initialize SSD
#device = torch.device('cuda:0')
#net = net.to('cuda:0')
net.load_state_dict(torch.load(args.weights))
transform = BaseTransform(net.size,
(104 / 256.0, 117 / 256.0, 123 / 256.0))
fps = FPS().start()
cv2_demo(net.eval(), transform)
# stop the timer and display FPS information
fps.stop()
print("[INFO] elasped time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))
# iter_size = accum_batch_size / batch_size
max_iter = args.iterations
weight_decay = 0.0005
stepvalues = (80000, 100000, 120000)
gamma = 0.1
momentum = 0.9
log_inter = 50
conf_thresh = 0.01 # default 0.01
nms_thresh = 0.45
if args.visdom:
import visdom
viz = visdom.Visdom(env="SSD OCR-MER")
ssd_net = build_ssd('train', version, 300, num_classes, conf_thresh=conf_thresh, nms_thresh=nms_thresh)
net = ssd_net
if args.cuda and ngpus > 1:
net = torch.nn.DataParallel(ssd_net)
# cudnn.benchmark = True
elif args.cuda:
net = net.cuda()
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
ssd_net.load_weights(args.resume)
else:
pretrained_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.base.load_state_dict(pretrained_weights)
type=int,
help='input size of network')
parser.add_argument('--load',
default='weights/trainer_004_2.5533.pth.tar',
help='input size of network')
parser.add_argument('--iid', default=None, help='single iid to test')
args = parser.parse_args()
print(args)
ds = DataSource()
test = ds.dataset(args.dataset)
if args.iid is not None: test.iids = [args.iid]
print(test.iids)
net = build_ssd('train', args.size, 6) # initialize SSD
if args.cuda:
net.cuda()
net.load_state_dict(torch.load(args.load))
else:
net.load_state_dict(
torch.load(args.load, map_location=lambda storage, loc: storage))
net.eval()
def gen_all(th, start=0):
print("start", start)
for iid in tqdm(test.iids[start:]):
ai, _ = test.aimg(iid)
mimg = ds.get_mimg(iid)
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_train = 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_train = VOCDetection(root=args.dataset_root,
image_sets=[('2012', 'trainval')],
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
dataset_val = VOCDetection(root=args.dataset_root,
image_sets=[('2007', 'val')],
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
dataset_test = VOCDetection(root=args.dataset_root,
image_sets=[('2007', 'test')],
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
# Take subsample of the datasets
frac = args.subset_size
data_sizes = np.array(
[len(dataset_train),
len(dataset_val),
len(dataset_test)])
data_sizes_sub = (data_sizes * frac).astype(int)
np.random.seed(10)
data_indices = [
np.random.choice(data_sizes[i], data_sizes_sub[i]) for i in range(3)
]
print("Train/Val/Test split " + str(data_sizes_sub[0]) + ':' +
str(data_sizes_sub[1]) + ':' + str(data_sizes_sub[2]))
if args.arch == 'ssd':
ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
net = ssd_net
elif args.arch == 'box_ssd':
ssd_net = build_ssd_box2('train', cfg['min_dim'], cfg['num_classes'])
net = ssd_net
elif args.arch == 'full_box_ssd':
ssd_net = build_ssd_full_box('train', cfg['min_dim'],
cfg['num_classes'])
net = ssd_net
else:
raise ("Incorrect Architecture chosen")
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:
# args.basenet == False
# pass
# vgg_weights = torch.load(args.save_folder + args.basenet)
# print('Loading base network...')
# ssd_net.vgg.load_state_dict(vgg_weights)
if args.cuda:
device = torch.device('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)
train_loss = []
valid_loss = []
train_loss_iter = []
epoch_time = []
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
loss_val = 0
print('Loading the dataset...')
epoch_size = len(dataset_train) // (args.batch_size / frac)
print('Training ' + args.arch + ' on:', dataset_train.name)
print('Using the specified args:')
print(args)
step_index = 0
if args.visdom:
vis_title = args.arch+' on ' + dataset_train.name + ' LR=' \
+ str(args.lr) + ' WD=' + str(args.weight_decay)
vis_legend_train = [
'Loc Loss Train', 'Conf Loss Train', 'Total Loss Train'
]
vis_legend_trainval = ['Train Loss', 'Val Loss']
iter_plot = create_vis_plot('Iteration', 'Loss', vis_title,
vis_legend_train)
epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title,
vis_legend_train)
loss_plot = viz.line(Y=torch.tensor([[0, 0]]).zero_(),
opts=dict(
title='Loss Tracker',
legend=['Training Loss', 'Validation Loss'],
xlabel='Iteration',
ylabel='Loss',
show_legend=True))
data_loader_train = data.DataLoader(dataset_train,
args.batch_size,
num_workers=args.num_workers,
sampler=SubsetRandomSampler(
data_indices[0]),
shuffle=False,
collate_fn=detection_collate,
pin_memory=True)
data_loader_val = data.DataLoader(dataset_val,
args.batch_size,
num_workers=args.num_workers,
sampler=SubsetRandomSampler(
data_indices[1]),
shuffle=False,
collate_fn=detection_collate,
pin_memory=True)
# data_loader_test = data.DataLoader(dataset_test, args.batch_size,
# num_workers=args.num_workers,
# sampler=SubsetRandomSampler(data_indices[2]),
# shuffle=False, collate_fn=detection_collate,
# pin_memory=True)
# import pdb; pdb.set_trace()
# mean = 0
# count = 0
# for t, (img,y) in enumerate(data_loader_train):
# mean += img.mean([0,2,3])
# print(img.mean([0,2,3]))
# count += 1
# if t % 10 == 0:
# print(mean/count)
# mean = mean/count
# print(mean)
# breakpoint()
# create batch iterator
batch_iterator = iter(data_loader_train)
for iteration in range(args.start_iter,
cfg['max_iter'] // 10): # //10 here to end early
if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
epoch += 1
print("Training Epoch number " + str(epoch))
epoch_start = time.time()
update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot, None,
'append', epoch_size)
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
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_train)
images, targets = next(batch_iterator)
if args.cuda:
images = images.to(device)
targets = [ann.to(device) for ann in targets]
# images = Variable(images.cuda())
# with torch.no_grad():
# targets = [Variable(ann.cuda()) for ann in targets]
# else:
# images = Variable(images)
# with torch.no_grad():
# targets = [Variable(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()
train_loss_iter.append(loss.item())
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Training Loss: %.4f ||' %
(loss.item()),
end=' ')
if args.visdom:
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(), 'weights/ssd300_COCO_' +
# repr(iteration) + '.pth')
# calculate val loss
#import pdb; pdb.set_trace()
if iteration != 0 and (iteration % epoch_size == 0):
net.eval()
loss_val = 0
with torch.no_grad():
for t, (images, targets) in enumerate(data_loader_val):
if args.cuda:
images = images.to(device)
targets = [ann.to(device) for ann in targets]
# images = images.cuda()
# targets = [Variable(ann.cuda()) for ann in targets]
out = net(images)
loss_l, loss_c = criterion(out, targets)
loss_val += loss_l + loss_c
loss_val /= len(data_loader_val)
# loc_loss += loss_l.item()
# conf_loss += loss_c.item()
print('iter ' + repr(iteration) + ' || Val Loss: %.4f ||' %
(loss_val.item()),
end=' ')
viz_tracker(loss_plot,
torch.tensor([[loss.item(),
loss_val.item()]]),
torch.tensor([[iteration - 1, iteration - 1]]))
#epoch += 1
# reset epoch loss counters
loss_l = 0
loss_c = 0
train_loss.append(loss.item())
valid_loss.append(loss_val.item())
#Epoch time
epoch_end = time.time()
epoch_time.append(epoch_end - epoch_start)
print("Epoch " + str(epoch) + " took " +
str(int(epoch_end - epoch_start)) + "secs to train")
suffix = args.arch + '_lr'+ str(args.lr) + '_wd' + str(args.weight_decay) \
+ '_sub'+ str(args.subset_size)
#Save resulting arrays so far every 10 epochs
if ((epoch) % 1 == 0):
with open('pkl_files/' + suffix + 'train_loss.pkl',
'wb') as handle:
pickle.dump(train_loss,
handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open('pkl_files/' + suffix + 'valid_loss.pkl',
'wb') as handle:
pickle.dump(valid_loss,
handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open('pkl_files/' + suffix + 'epoch_time.pkl',
'wb') as handle:
pickle.dump(epoch_time,
handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open('pkl_files/' + suffix + 'train_loss_iter.pkl',
'wb') as handle:
pickle.dump(train_loss_iter,
handle,
protocol=pickle.HIGHEST_PROTOCOL)
state = {
'epoch': epoch,
'state_dict': ssd_net.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(
state, args.save_folder + '' + str(args.subset_size) +
args.dataset + suffix + '.pth')
means = (104, 117, 123) # only support voc now
num_classes = 3 #english, not english, na
batch_size = args.batch_size
accum_batch_size = 32
iter_size = accum_batch_size / batch_size
max_iter = 400000
weight_decay = 0.0005
stepvalues = (280000, 360000, 400000)
gamma = 0.1
momentum = 0.9
if args.visdom:
import visdom
viz = visdom.Visdom()
ssd_net = build_ssd('train', ssd_dim, num_classes)
net = ssd_net
if args.cuda:
net = torch.nn.DataParallel(ssd_net)
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 args.cuda:
net.cuda()
def train():
cfg = balloon # config.py中的voc字典
dataset = BALLOONDetection(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: # 加载vgg16的预训练模型
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_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('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)
print('dataset:{}'.format(len(dataset)))
print('dataloader:{}'.format(len(data_loader)))
# create batch iterator
batch_iterator = iter(data_loader)
print('batch iterator:{}'.format(len(batch_iterator)))
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']: # 如果到了lr变化的节点step
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
try:
images, targets = next(batch_iterator) # images = [32, 3, 300, 300]
# print(images.shape, end="")
# for i in targets:
# print(i.shape, end="")
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.item()
conf_loss += loss_c.item()
if iteration % 10 == 0: # 每500显示一次loss
print('timer: %.4f sec.' % (t1 - t0))
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')
if iteration != 0 and iteration % 5000 == 0: # 每5000次保存一次模型
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(), 'weights/ssd300_' + args.dataset + '_' +
repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(), args.save_folder + '' + args.dataset + '.pth') # 最后保存一次模型
def demo_plt(img_id=0):
net = build_ssd('test', 512, 5) # initialize SSD
print(net)
net.load_weights(
'/media/sunwl/Datum/Projects/GraduationProject/Fused_sum_SSD_VHR_512_conv3_3/weights/v2_rosd.pth'
)
testset = ROSDDetection(ROSDroot, ['test'], None, AnnotationTransform_ROSD)
image = testset.pull_image(img_id)
# image = cv2.imread('demos/047.jpg')
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# View the sampled input image before transform
plt.figure(figsize=(10, 10))
plt.imshow(rgb_image)
x = cv2.resize(rgb_image, (512, 512)).astype(np.float32)
x -= (104.0, 117.0, 123.0)
x = x.astype(np.float32)
x = x[:, :, ::-1].copy()
x = torch.from_numpy(x).permute(2, 0, 1)
xx = Variable(x.unsqueeze(0)) # wrap tensor in Variable
if torch.cuda.is_available():
xx = xx.cuda()
y = net(xx)
#
plt.figure(figsize=(10, 10))
colors = plt.cm.hsv(np.linspace(0, 1, 5)).tolist()
plt.imshow(rgb_image.astype(np.uint8)) # plot the image for matplotlib
currentAxis = plt.gca()
detections = y.data
# scale each detection back up to the image
scale = torch.Tensor(rgb_image.shape[1::-1]).repeat(2)
for i in range(detections.size(1)):
j = 0
while detections[0, i, j, 0] >= 0.6:
score = detections[0, i, j, 0]
label_name = labels[i - 1]
display_txt = '%s: %.2f' % (label_name, score)
pt = (detections[0, i, j, 1:] * scale).cpu().numpy()
color = colors[i]
coords = (pt[0], pt[1]), pt[2] - pt[0] + 1, pt[3] - pt[1] + 1
currentAxis.add_patch(
plt.Rectangle(*coords,
fill=False,
edgecolor=color,
linewidth=2))
currentAxis.text(pt[0],
pt[1],
display_txt,
bbox={
'facecolor': color,
'alpha': 0.5
})
j += 1
plt.savefig(
'/media/sunwl/Datum/Projects/GraduationProject/Fused_sum_SSD_VHR_512_conv3_3/outputs/rosd_{:03}.png'
.format(img_id))
plt.xticks([])
plt.yticks([])
plt.show()
def train():
cfg = voc
dataset = MyDetection(root=os.path.join(HOME, "workspace/Data/"),
datasetfile='8000list',
transform=SSDAugmentation(cfg['min_dim'], MEANS))
if args.visdom:
import visdom
viz = visdom.Visdom()
ssd_net = build_ssd('train', cfg['min_dim'], len(My_CLASSES) + 1)
net = ssd_net
''' no multi-gpu
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:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_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('Training SSD:')
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)
batch_iterator = iter(cycle(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)
images = Variable(images.cuda())
targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
# forward
t0 = time.time()
out = net(images)
# print(out)
# 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]
if iteration % 1 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' %
(loss.data[0]),
end=' ')
#print("\r +++++ \r loss:", loss.data[0], loss_l.data[0], loss_c.data[0])
if args.visdom:
update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(),
'weights/ssd300_' + repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
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))
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:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_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)
now_time = (datetime.datetime.now() +
datetime.timedelta(hours=8)).strftime('%Y-%m-%d-%H-%M')
writer = SummaryWriter("runs/" + now_time[5:])
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
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)
loc_loss_epoch = 0.0
conf_loss_epoch = 0.0
total_loss_epoch = 0.0
# 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)
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
if args.cuda:
images = images.cuda()
for i in range(len(targets)):
targets[i] = targets[i].float().cuda()
else:
images = images
targets = 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()
loc_loss_epoch += loss_l.item()
conf_loss_epoch += loss_c.item()
total_loss_epoch += loss.item()
if (iteration + 1) % epoch_size == 0:
writer.add_scalar("loss/loc_loss", loc_loss_epoch / len(dataset),
(iteration + 1) // epoch_size)
writer.add_scalar("loss/conf_loss", conf_loss_epoch / len(dataset),
(iteration + 1) // epoch_size)
writer.add_scalar("loss/total_loss",
total_loss_epoch / len(dataset),
(iteration + 1) // epoch_size)
loc_loss_epoch = 0.0
conf_loss_epoch = 0.0
total_loss_epoch = 0.0
if (iteration + 1) % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter [' + repr(iteration + 1) + "/" + str(cfg["max_iter"]) +
'] Loss: %.4f ||' % (loss.item()),
end=' ')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(),
'weights/ssd300_voc_pf_' + repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
num_classes = len(VHR_CLASSES) + 1
batch_size = args.batch_size
accum_batch_size = 1
iter_size = accum_batch_size / batch_size
max_iter = 120000
weight_decay = args.weight_decay
stepvalues = (80000, 100000, 120000)
gamma = args.gamma
momentum = args.momentum
if args.visdom:
import visdom
viz = visdom.Visdom()
ssd_net = build_ssd('train', 512, num_classes)
net = ssd_net
if args.cuda:
net = torch.nn.DataParallel(ssd_net)
cudnn.benchmark = True
def xavier(param):
init.xavier_uniform(param)
def weights_init(model):
if isinstance(model, list):
for m in model:
weights_init(m)
def train():
cfg = voc
dataset = VOCDetection(root=VOC_ROOT,
transform=SSDAugmentation(cfg['min_dim'], MEANS))
ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = ssd_net.to(device)
vgg_weights = torch.load('weights/vgg16.pth')
ssd_net.vgg.load_state_dict(vgg_weights)
if gpu:
net = torch.nn.DataParallel(ssd_net)
cudnn.benchmark = True
ssd_net.extras.apply(weights_init)
ssd_net.loc.apply(weights_init)
ssd_net.conf.apply(weights_init)
# 損失関数の設定
criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
False, gpu)
# 最適化パラメータの設定
optimizer = optim.SGD(net.parameters(),
lr=lr_,
momentum=0.9,
weight_decay=weight_decay)
# 訓練モード
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // batch_size
print('dataset_size', len(dataset))
print('epoch_size', epoch_size)
print('Training SSD on:', dataset.name)
step_index = 0
data_loader = data.DataLoader(dataset,
batch_size,
num_workers=4,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
# 学習の開始
batch_iterator = None
# iterationでループして、cfg['max_iter']まで学習する
for iteration in range(0, 1000):
# 学習開始時または1epoch終了後にdata_loaderから訓練データをロードする
if (not batch_iterator) or (iteration % epoch_size == 0):
batch_iterator = iter(data_loader)
loc_loss = 0
conf_loss = 0
epoch += 1
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, gamma_, step_index)
# load train data
# バッチサイズ分の訓練データをload
images, targets = next(batch_iterator)
# 画像をGPUに転送
images = images.to(device)
# アノテーションをGPUに転送
targets = [ann.to(device) for ann in targets]
# forward
t0 = time.time()
# 順伝播の計算
out = net(images)
# 勾配の初期化
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=' ')
# 学習済みモデルの保存
torch.save(ssd_net.state_dict(), 'weights/ssd.pth')
parser.add_argument('--classes',
default=['horse'],
nargs='+',
help='The class, that shal be recognised in the image')
args = parser.parse_args()
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
lr = args.lr
torch.set_default_tensor_type('torch.cuda.FloatTensor')
img_size = 300
#criterion = MultiBoxLoss(21, 0.5, True, 0, True, 3, 0.5, False, True)
net = build_ssd('train', img_size, 21) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
criterion = MultiBoxLoss(21, 0.5, True, 0, True, 3, 0.5, False, True)
prep = transforms.Compose([
transforms.Scale((img_size, img_size)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
postpa = transforms.Compose([
transforms.Normalize(mean=[0, 0, 0], std=[1 / 0.229, 1 / 0.224,
1 / 0.225]),
transforms.Normalize(mean=[-0.485, -0.456, -0.406], std=[1, 1, 1]),
])
frame_t = transform(frame)[0] # We apply the transformation to our frame.
x = torch.from_numpy(frame_t).permute(2, 0, 1) # We convert the frame into a torch tensor.
x = Variable(x.unsqueeze(0)) # We add a fake dimension corresponding to the batch.
y = net(x) # We feed the neural network ssd with the image and we get the output y.
detections = y.data # We create the detections tensor contained in the output y.
scale = torch.Tensor([width, height, width, height]) # We create a tensor object of dimensions [width, height, width, height].
for i in range(detections.size(1)): # For every class:
j = 0 # We initialize the loop variable j that will correspond to the occurrences of the class.
while detections[0, i, j, 0] >= 0.6: # We take into account all the occurrences j of the class i that have a matching score larger than 0.6.
pt = (detections[0, i, j, 1:] * scale).numpy() # We get the coordinates of the points at the upper left and the lower right of the detector rectangle.
cv2.rectangle(frame, (int(pt[0]), int(pt[1])), (int(pt[2]), int(pt[3])), (255, 0, 0), 2) # We draw a rectangle around the detected object.
cv2.putText(frame, labelmap[i - 1], (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 2, cv2.LINE_AA) # We put the label of the class right above the rectangle.
j += 1 # We increment j to get to the next occurrence.
return frame # We return the original frame with the detector rectangle and the label around the detected object.
# Creating the SSD neural network
net = build_ssd('test') # We create an object that is our neural network ssd.
net.load_state_dict(torch.load('ssd300_mAP_77.43_v2.pth', map_location = lambda storage, loc: storage)) # We get the weights of the neural network from another one that is pretrained (ssd300_mAP_77.43_v2.pth).
# Creating the transformation
transform = BaseTransform(net.size, (104/256.0, 117/256.0, 123/256.0)) # We create an object of the BaseTransform class, a class that will do the required transformations so that the image can be the input of the neural network.
# Doing some Object Detection on a video
reader = imageio.get_reader('funny_dog.mp4') # We open the video.
fps = reader.get_meta_data()['fps'] # We get the fps frequence (frames per second).
writer = imageio.get_writer('output.mp4', fps = fps) # We create an output video with this same fps frequence.
for i, frame in enumerate(reader): # We iterate on the frames of the output video:
frame = detect(frame, net.eval(), transform) # We call our detect function (defined above) to detect the object on the frame.
writer.append_data(frame) # We add the next frame in the output video.
print(i) # We print the number of the processed frame.
writer.close() # We close the process that handles the creation of the output video.
module_path = os.path.abspath(os.path.join('...'))
if module_path not in sys.path:
sys.path.append(module_path)
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
import numpy as np
import cv2
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
from ssd import build_ssd
net = build_ssd('train', 300, 21) # initialize SSD
net.load_weights('weights/ssd300_mAP_77.43_v2.pth')
# image = cv2.imread('./data/example.jpg', cv2.IMREAD_COLOR) # uncomment if dataset not downloaded
from matplotlib import pyplot as plt
from data import VOCDetection, VOC_ROOT, VOCAnnotationTransform
# here we specify year (07 or 12) and dataset ('test', 'val', 'train')
testset = VOCDetection(VOC_ROOT, [('2007', 'val')], None,
VOCAnnotationTransform())
img_id = 60
image = testset.pull_image(img_id)
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# View the sampled input image before transform
plt.figure(figsize=(10, 10))
plt.imshow(rgb_image)
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
print('Evaluating detections')
evaluate_detections(all_boxes, output_dir, dataset)
def evaluate_detections(box_list, output_dir, dataset):
write_voc_results_file(box_list, dataset)
do_python_eval(output_dir)
if __name__ == '__main__':
# load net
num_classes = len(labelmap) + 1 # +1 for background
net = build_ssd('test', 300, num_classes) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
# load data
dataset = VOCDetection(args.voc_root, [('2007', set_type)],
BaseTransform(300, dataset_mean),
VOCAnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(args.save_folder, net, args.cuda, dataset,
BaseTransform(net.size, dataset_mean), args.top_k, 300,
thresh=args.confidence_threshold)
def train():
# loading dataset and set the default configuraitons
cfg = mic
dataset = MICDetection(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:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_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)
"""
num_classes, overlap_thresh, prior_for_matching,
bkg_label, neg_mining, neg_pos, neg_overlap, encode_target,
use_gpu=True
"""
criterion = MultiBoxLoss(cfg['num_classes'], 0.1, 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('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
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.item()
conf_loss += loss_c.item()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' %
(loss.data.item()),
end=' ')
print('->> conf loss:{:.4f} || loc loss:{:.4f}'.format(
loss_c.data.item(), loss_l.data.item()))
print('->>lr:{:.6f}'.format(optimizer.param_groups[0]['lr']))
if args.visdom:
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(),
'weights/ssd300_MIC_' + repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
def eval():
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 == 'wamv':
cfg = wamv
dataset = wamvDetection(root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif args.dataset == 'subt':
cfg = subt
dataset = subtDetection(root=args.dataset_root,
image_sets=[('test')],
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
ssd_net = build_ssd('test', 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:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_weights)
if args.cuda:
net = net.cuda()
net.eval()
# loss counters
testset = subtDetection(subt_ROOT, [('train')], None,
subtAnnotationTransform())
Tensor = torch.cuda.FloatTensor if args.cuda else torch.FloatTensor
imgs_list = [] # Stores image paths
img_detections = [] # Stores detections for each image index
result_list = [0, 0, 0, 0] # TP TN FP FN
all_data = 0
conf_thres = 0
iou_thres = 0.5
print("\nPerforming object detection:")
scale = torch.Tensor((640, 480)).repeat(2)
for batch_i in range(len(testset.ids)):
image = testset.pull_image(batch_i)
x = cv2.resize(image, (300, 300)).astype(np.float32)
x -= (104.0, 117.0, 123.0)
x = x.astype(np.float32)
x = x[:, :, ::-1].copy()
x = torch.from_numpy(x).permute(2, 0, 1)
xx = Variable(x.unsqueeze(0))
if torch.cuda.is_available():
xx = xx.cuda()
with torch.no_grad():
detections = net(xx).data
# Save image and detections
# imgs_list.extend(imgs)
# img_detections.extend(detections)
targets = testset.pull_anno(batch_i)[1]
for label in targets:
all_data += 1
pt_label = label[0:4]
mask_label = np.zeros((640, 480), np.float64)
mask_label[int(pt_label[0]):int(pt_label[2]),
int(pt_label[1]):int(pt_label[3])] = 1
_bool = False
for i in range(detections.size(1)):
j = 0
while detections[0, i, j, 0] > conf_thres:
score = detections[0, i, j, 0]
pt = (detections[0, i, j, 1:] * scale).cpu().numpy()
mask_predict = np.zeros((640, 480), np.float64)
mask_predict[int(pt[0]):int(pt[2]),
int(pt[1]):int(pt[3])] = 1
_and = sum(sum(np.logical_and(mask_label, mask_predict)))
_or = sum(sum(np.logical_or(mask_label, mask_predict)))
_iou = float(_and / _or)
# print (_iou)
if score >= conf_thres:
if (i - 1
) == label[4] and not _bool and _iou >= iou_thres:
_bool = True
result_list[0] += 1
elif (i - 1) != label[4] or _iou <= iou_thres:
print(i - 1, label[4], _iou)
result_list[2] += 1
j += 1
if not _bool:
result_list[1] += 1
print(result_list)
print(all_data)
print("Recall : ", float(result_list[0] / all_data))
print("Precision : ",
float(result_list[0] / (result_list[0] + result_list[2])))
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
print('Evaluating detections')
evaluate_detections(all_boxes, output_dir, dataset)
return output_dir
def evaluate_detections(box_list, output_dir, dataset):
write_voc_results_file(box_list, dataset)
do_python_eval(output_dir)
if __name__ == '__main__':
# load net
num_classes = len(labelmap) + 1 # +1 for background
net = build_ssd('test', args.resize_size, num_classes,
configure=cfg) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
# load data
dataset = VOCDetection(args.voc_root, [('2007', set_type)],
BaseTransform(args.resize_size, dataset_mean),
VOCAnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
cachedir = os.path.join(devkit_path, 'annotations_cache')
if os.path.exists(cachedir):
shutil.rmtree(cachedir)
default='weights/day1_001_2.4164.pth.tar',
help='input size of network')
parser.add_argument('--iid', default=None, help='single iid to test')
parser.add_argument('--name', default="day1", help='experiment name')
parser.add_argument('--plot', default=True, type=bool, help='experiment name')
parser.add_argument('--mask', default=True, type=bool, help='use mask')
args = parser.parse_args()
datasource = DataSource()
test = datasource.dataset(args.dataset)
test.shuffle()
test.take(start=args.start, end=args.end)
model = build_ssd('test',
args.size,
args.classes,
load=args.load,
cuda=args.cuda)
datasource.mkpath('anns_test', test=args.name, dataset=args.dataset)
if args.plot:
datasource.mkpath('plot_test', test=args.name, dataset=args.dataset)
iids = test.iids
ret = np.loadtxt("/home/arun/Desktop/submission.csv",
delimiter=',',
skiprows=1).astype('int32')
print(ret[:10, 0])
iids = ret[:, 0]
iids = iids[86:]
if args.iid is not None: iids = [args.iid]
for iid in tqdm(iids):
aimg, iid = test.aimg(iid)
cv2.imshow("frame", frame)
if key == 27: # exit
stream.stop()
break
if __name__ == "__main__":
import sys
from os import path
sys.path.append(path.dirname(path.dirname(path.abspath(__file__))))
from data import BaseTransform, VOC_CLASSES as labelmap
from ssd import build_ssd
net = build_ssd("test", 300, 21) # initialize SSD
# if args.cuda:
# net = torch.nn.DataParallel(net)
# cudnn.benchmark = True
# state_dict = torch.load(args.weights)
# from collections import OrderedDict
#
# new_state_dict = OrderedDict()
# for k, v in state_dict.items():
# name = 'module.' + k # remove `module.`
# new_state_dict[name] = v
# net.load_state_dict(new_state_dict)
# net.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(args.weights)['state_dict'].items()})
net.load_state_dict(torch.load(args.weights))
# net.load_state_dict(torch.load(args.weights,map_location='cuda'))
transform = BaseTransform(net.size,
means = (104, 117, 123) # only support voc now
num_classes = len(VOC_CLASSES) + 1
batch_size = args.batch_size
accum_batch_size = 32
iter_size = accum_batch_size / batch_size
max_iter = 120000
weight_decay = 0.0005
stepvalues = (80000, 100000, 120000)
gamma = 0.1
momentum = 0.9
if args.visdom:
import visdom
viz = visdom.Visdom()
ssd_net = build_ssd('train', 300, 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:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_weights)
if args.cuda:
# torch.set_default_tensor_type('torch.cuda.FloatTensor')
# if not cuda:
# print("WARNING: It looks like you have a CUDA device, but aren't " +
# "using CUDA.\nRun with --cuda for optimal training speed.")
# torch.set_default_tensor_type('torch.FloatTensor')
# else:
# torch.set_default_tensor_type('torch.FloatTensor')
if not os.path.exists(save_folder):
os.mkdir(save_folder)
cfg = voc
dataset = VOCDetection(root='data/',
transform=SSDAugmentation(cfg['min_dim'], MEANS))
ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
net = ssd_net
if cuda:
net = torch.nn.DataParallel(ssd_net)
cudnn.benchmark = True
if resume:
print('Resuming training, loading {}...'.format(resume))
ssd_net.load_weights(resume)
else:
vgg_weights = torch.load(save_folder + basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_weights)
if cuda:
def train():
cfg = voc # voc是一个字典 里面包括网络的一系列参数信息
dataset = VOCDetection( # 是一个VOC数据的类
root=args.dataset_root, # 数据集的根目录
transform=SSDAugmentation(
cfg['min_dim'],
MEANS)) # 图片的预处理方法(输入图片的尺寸和均值) 原本类中定义为None 后面的MEANS我人为可以删除
if args.visdom: # 这里是可视化工具，不用管###################
import visdom
viz = visdom.Visdom()
ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
# 阶段【train or test】 输入图片尺寸大小 类别数
# build_ssd是一个放在ssd.py的函数
# return是一个类的对象，也就是class SSD(nn.Module)，ssd_net也就是SSD类的一个对象
# ssd_net拥有所有class SSD继承于nn.Module以及作者增加方法的所有属性
# 在SSD这个类中就定义了网络的base部分（修改全连接层后的VGG16）和extras部分（论文作者加入的多尺度feature map）和head部分
# 对选定的6个尺度下的feature map进行卷积操作得到的每个default box 的每一个分类类别的confidence以及位置坐标的信息
net = ssd_net # 到这里class类SSD只完成了__init__()并没有执行__forward__() net是一个类
if args.cuda: # 是否将模型放到多个个GPU上运行{我认为在我的任务中不要放在多线程GPU中}
net = torch.nn.DataParallel(ssd_net)
cudnn.benchmark = True
if args.resume: # 【resume】的默认值是None,表示不是接着某个断点来继续训练这个模型 【其实checkpoint里面最好还要加上优化器的保存】
# 【model_state_dict,optimizer_state_dict,epoch】 见深度之眼
print('Resuming training, loading {}...'.format(args.resume))
ssd_net.load_weights(args.resume)
else: # 那么就从weights文件夹下面直接加载预训练好vgg基础网络预训练权重
vgg_weights = torch.load(args.save_folder +
args.basenet) # 整个ssd_net中vgg基础网络的权重
print('Loading base network...')
ssd_net.vgg.load_state_dict(
vgg_weights
) # 只在整个ssd_net中的vgg模块中加载预训练好的权重，其余的extra，特征融合，CBAM模块没有加载预训练权重
if args.cuda: # 将模型结构放在GPU上训练
net = net.cuda()
if not args.resume: # ######################################################################
print('Initializing weights...'
) # 如果不是接着某个断点接着训练，那么其余extras loc con都会xavier方法初始化
# initialize newly added layers' weights with xavier method
ssd_net.extras.apply(
weights_init) # extras 模块由 xavier 方法默认初始化data和bias
ssd_net.loc.apply(weights_init) # loc 模块由 xavier 方法默认初始化data和bias
ssd_net.conf.apply(weights_init) # conf 模块由 xavier 方法默认初始化data和bias
# 【优化器】net.parameters()是网络结构中的参数，学习率，动量，权重衰减率
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# 定义损失函数部分【MultiBoxesLoss是一个类用于计算网络的损失，criterion是一个对象】
# 【损失函数】 关键！！！ criterion是个nn.Moudule的形式 里面包括两部分loss_c 和 loss_l
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 # 每个epoch中有多少个batch
print('Training SSD on:', dataset.name)
print('Using the specified args:')
print(args) # 讲设定的参数打印出来
step_index = 0
# 可视化部分
if args.visdom: # 默认值为False
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, # 默认值我修改成了0
shuffle=True,
collate_fn=detection_collate, # collate_fn将一个batch_size数目的图片进行合并成batch
pin_memory=True)
batch_iterator = iter(data_loader) # batch迭代器 依次迭代batch
for iteration in range(args.start_iter, cfg['max_iter']): # 由最大迭代次数来迭代训练
if args.visdom and iteration != 0 and (
iteration % epoch_size == 0): # 因为args.visdom一直设置为False因此没有被调用
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']: # 通过多少次epoch调节一次学习率
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
try:
images, targets = next(batch_iterator)
# targets 和image都是读取的训练数据
except StopIteration:
bath_iterator = iter(data_loader)
images, targets = next(bath_iterator)
# images=【batch_size,3,300,300】
# targets=【batch_size,num_object,5】
# num_object代表一张图里面有几个ground truth，5代表四个位置信息和一个label
if args.cuda: # 将数据放在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是netforward的输出：是个元组，里面包括3个部分[loc conf priors]
out = net(images)
# ## backprop 优化器梯度清零
optimizer.zero_grad()
# ## criterion是nn.Module形式，下面是调用它的forward模式【重点看，里面包括难例挖掘的内容】
# ###################################【【【训练阶段的损失！！！】】】######################################
# ##输入参数1：网络结构net输出的out:[loc conf priors]
# ##输入参数2：targets:真实目标的位置标签值
loss_l, loss_c = criterion(
out, targets
) # criterion就是MultiBoxLoss类定义的对象，forward前传播返回的结果是【loss_l, loss_c】
loss = loss_l + loss_c # 总loss
loss.backward()
optimizer.step()
t1 = time.time()
# 下面两行好像没有使用
loc_loss += loss_l.data # ###到底是改成item()还是data
conf_loss += loss_c.data # ###到底是改成item()还是data
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % loss.data,
end=' ') # 到底是改成item()还是data
if args.visdom:
update_vis_plot(iteration, loss_l.data, loss_c.data, iter_plot,
epoch_plot, 'append')
if iteration != 0 and iteration % 2000 == 0:
# 迭代多少次保存一次模型。 在尝试阶段，为了节省时间，建议将根据迭代次数保存模型的参数调低，例如调节到500
print('Saving state, iter:', iteration) # 保存的checkpoint
torch.save(ssd_net.state_dict(), 'weights/ssd300_VOC_' +
repr(iteration) + '.pth') # 保存模型的路径
torch.save(ssd_net.state_dict(), args.save_folder + '' + args.dataset +
'.pth') # 最后的保存：不是保存整个模型，只是保存了参数
def train():
writer = SummaryWriter()
cfg = voc
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'],
MEANS))
ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
net = ssd_net
if args.visdom:
import visdom
viz = visdom.Visdom()
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:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_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.AdamW(net.parameters(), lr=args.lr)
# 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('Training SSD on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
data_loader = torch.utils.data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
# create batch iterator
for epoch in range(10000):
for iteration, data_ in enumerate(data_loader):
# if iteration in cfg['lr_steps']:
# step_index += 1
# adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
images, targets = data_
if args.cuda:
images = images.cuda()
targets = [ann.cuda().detach() for ann in targets]
else:
images = Variable(images)
targets = [ann.detach() for ann in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
loc_loss += loss_l.item()
conf_loss += loss_c.item()
# torch.save(ssd_net.state_dict(),
# args.save_folder + '' + args.dataset + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + 'VOC_weighted' + '.pth')
writer.add_scalar("epoch location loss", conf_loss, epoch)
writer.add_scalar("epoch classification loss", loc_loss, epoch)
writer.add_scalar("epoch total loss", loss.item(), epoch)
writer.flush()
loc_loss = 0
conf_loss = 0
writer.close()
import sys
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
import numpy as np
import cv2
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
from ssd import build_ssd
# SSDネットワークの定義と重みファイルのロード
net = build_ssd('test', 300, 21)
#net.load_weights('./weights/BCCD.pth')
net.load_weights('./weights/enseble_stars_ssd.pth')
from matplotlib import pyplot as plt
from data import VOCDetection, VOC_ROOT, VOCAnnotationTransform
# BCCD_test 読み込み
#testset = VOCDetection(VOC_ROOT, [('BCCD', 'test')], None, VOCAnnotationTransform())
testset = VOCDetection(VOC_ROOT, [('enseble_stars_ssd', 'test')], None, VOCAnnotationTransform())
img_id = 25
image = testset.pull_image(img_id)
# テスト画像の表示
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(10,10))
plt.imshow(rgb_image)
plt.show()
encoding='utf-8')
class_name_file = "/home/hyer/datasets/chinese/voc_data_200/code_200.txt"
#
class_names = []
with open(class_name_file, "r") as f:
label_data = f.readlines()
for li in label_data:
name = li.strip()
class_names.append(uni2chi(name))
# cudnn.benchmark = True # maybe faster
net = build_ssd('test',
version,
size=300,
num_classes=num_cls,
bkg_label=0,
top_k=100,
conf_thresh=0.1,
nms_thresh=0.45) # initialize SSD
net.load_weights(
'/home/hyer/workspace/algo/Detection/SSD/ssd.pytorch/models/VOC_mobile_300_19105321_5000-cls201.pth'
)
net.eval() # important!!!
if use_cuda:
net = net.cuda()
from matplotlib import pyplot as plt
_t = {'im_detect': Timer(), 'misc': Timer()}
# -*- coding: UTF-8 -*-
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
print('Evaluating detections')
evaluate_detections(all_boxes, output_dir, dataset)
def evaluate_detections(box_list, output_dir, dataset):
write_yto_results_file(box_list, dataset)
do_python_eval(output_dir)
if __name__ == '__main__':
# load net
num_classes = len(YTO_CLASSES) + 1 # +1 background
net = build_ssd('test', 300, num_classes) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
# load data
dataset = YTODetection(args.yto_root, [(set_type)],
BaseTransform(300, dataset_mean),
yto_AnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(args.save_folder,
net,
args.cuda,
dataset,
parser.add_argument('--exp', default="day1", help='exp name')
parser.add_argument('--epochs', default=80, type=int, help='exp name')
parser.add_argument('--fcount', default=1, type=int, help='filter count')
parser.add_argument('--aug', default=True, type=bool, help='use augmentations')
parser.add_argument('--limit',
default=None,
type=int,
help='limit the dataset')
args = parser.parse_args()
print(args)
cfg = (v1, v2)[args.version == 'v2']
torch.set_default_tensor_type('torch.FloatTensor')
#print(torch.zeros(1))
#model = Network()
print("building ssd")
model = build_ssd('train', args.dim, args.num_classes)
print("building ssd done")
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
model.vgg.load_state_dict(vgg_weights)
def xavier(param):
init.xavier_uniform(param)
def weights_init(m):
if isinstance(m, nn.Conv2d):
xavier(m.weight.data)
m.bias.data.zero_()
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))
if args.visdom:
import visdom
#LGQ global visdom
global viz
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:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_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('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)
#LGQ modified
# solve the problem of iteration is used up
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
global update_iterater
update_iterater = 1
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]
#LGQ modified
loc_loss += loss_l.data
conf_loss += loss_c.data
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
#LGQ modified
#print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' %
(loss.data),
end=' ')
if args.visdom:
#LGQ update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],
update_vis_plot(iteration, loss_l.data, loss_c.data, iter_plot,
epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(),
'weights/ssd300_COCO_' + repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
#import pdb;pdb.set_trace()
print(args)
if torch.cuda.is_available():
if args.cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
if not args.cuda:
print(
"WARNING: It looks like you have a CUDA device, but aren't using \
CUDA. Run with --cuda for optimal eval speed.")
torch.set_default_tensor_type('torch.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
cfg = casiaship
# load net
net = build_ssd('test', cfg, cfg['min_dim'], cfg['num_classes'])
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
#imagesetdir
imagedir = os.path.join(args.data_dir, 'image')
# imagedir=args.data_dir
annotdir = os.path.join(args.data_dir, 'label')
imglist = get_infer_imagelist(imagedir, args.test_images)
# import pdb;pdb.set_trace()
#genarate savefolder
det_path, imagenames, vis_dir = mksavetree(args.save_folder)
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))
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:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_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('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)
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]
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')
if args.visdom:
update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(), 'weights/ssd300_COCO_' +
repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
if key2 == ord('p'): # resume
break
cv2.imshow('frame', frame)
if key == 27: # exit
break
if __name__ == '__main__':
import sys
from os import path
sys.path.append(path.dirname(path.dirname(path.abspath(__file__))))
from data import BaseTransform, VOC_CLASSES as labelmap
from ssd import build_ssd
net = build_ssd('test', 300, 21) # initialize SSD
net.load_state_dict(torch.load(args.weights))
transform = BaseTransform(net.size, (104/256.0, 117/256.0, 123/256.0))
fps = FPS().start()
cv2_demo(net.eval(), transform)
# stop the timer and display FPS information
fps.stop()
print("[INFO] elasped time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))
# cleanup
cv2.destroyAllWindows()
stream.stop()