def get_num_train_image(self):
if op.isfile(self.get_data('trainX')):
if op.isfile(self.get_shuffle_file('train')):
return len(load_list_file(self.get_shuffle_file('train')))
else:
return 0
else:
return len(load_list_file(op.join(self._data_root, 'train.lineidx')))
def get_train_tsvs(self, t=None):
if op.isfile(self.get_data('train', t)):
return [self.get_data('train', t)]
trainx_file = op.join(self._data_root, 'trainX.tsv')
if not op.isfile(trainx_file):
return []
train_x = load_list_file(trainx_file)
if t is None:
return train_x
elif t =='label':
if op.isfile(self.get_data('trainX', 'label')):
return load_list_file(self.get_data('trainX', 'label'))
else:
files = [op.splitext(f)[0] + '.label.tsv' for f in train_x]
return files
def get_all_data_info2(name=None):
if name is None:
return sorted(os.listdir('./data'))
else:
dataset = TSVDataset(name)
if not op.isfile(dataset.get_labelmap_file()):
return []
global_labelmap = None
labels = dataset.load_labelmap()
# here we assume the composite dataset has only one version
valid_split_versions = []
if len(dataset.get_train_tsvs()) > 1:
global_labelmap = dataset.load_labelmap() if global_labelmap is \
None else global_labelmap
valid_split_versions.append(('train', 0, global_labelmap))
splits = ['trainval', 'test']
else:
splits = ['train', 'trainval', 'test']
for split in splits:
v = 0
while True:
if not op.isfile(dataset.get_data(split, 'label', v)):
break
valid_split_versions.append((split, v, load_list_file(dataset.get_data(split,
'labelmap', v))))
v = v + 1
name_splits_labels = [(name, valid_split_versions)]
return name_splits_labels
def load_inverted_label(self, split, version=None, label=None):
fname = self.get_data(split, 'inverted.label', version)
if not op.isfile(fname):
return {}
elif label is None:
rows = tsv_reader(fname)
result = {}
for row in rows:
assert row[0] not in result
assert len(row) == 2
ss = row[1].split(' ')
if len(ss) == 1 and ss[0] == '':
result[row[0]] = []
else:
result[row[0]] = map(int, ss)
return result
else:
all_label = load_list_file(self.get_data(split, 'labelmap', version))
result = {}
idx = all_label.index(label)
row = TSVFile(fname).seek(idx)
assert row[0] == label
ss = row[1].split(' ')
if len(ss) == 1 and ss[0] == '':
result[row[0]] = []
else:
result[row[0]] = map(int, ss)
return result
def tsv_shuffle_reader(tsv_file):
logging.warn('deprecated: using TSVFile to randomly seek')
lineidx_file = op.splitext(tsv_file)[0] + '.lineidx'
lineidx = load_list_file(lineidx_file)
random.shuffle(lineidx)
with open(tsv_file, 'r') as fp:
for l in lineidx:
fp.seek(int(float(l)))
yield [x.strip() for x in fp.readline().split('\t')]
def __init__(self, tsv_file, labelmap, transform=None):
from tsv_io import TSVFile
from qd_common import load_list_file
self._tsv = TSVFile(tsv_file)
self._label_to_idx = {
l: i
for i, l in enumerate(load_list_file(labelmap))
}
self.transform = transform
self.ids = None
def iter_data(self, split, t=None, version=None):
if split == 'train' and op.isfile(self.get_data('trainX')):
assert version is None
train_files = load_list_file(self.get_data('trainX', t))
train_tsvs = [TSVFile(f) for f in train_files]
train_label_files = load_list_file(self.get_data('trainX',
'label'))
train_label_tsvs = [TSVFile(f) for f in train_label_files]
shuffle_file = self.get_shuffle_file('train')
shuffle_tsv_rows = tsv_reader(shuffle_file)
for idx_source, idx_row in shuffle_tsv_rows:
idx_source, idx_row = int(idx_source), int(idx_row)
data_row = train_tsvs[idx_source].seek(idx_row)
label_row = train_label_tsvs[idx_source].seek(idx_row)
assert label_row[0] == data_row[0]
yield label_row[0], label_row[1], data_row[-1]
else:
if not op.isfile(self.get_data(split, t, version)):
return
for row in tsv_reader(self.get_data(split, t, version)):
yield row
def load_label_parent(label_tree_file):
buf = read_to_buffer(label_tree_file).split('\n')
label_map = [b.split(' ') for b in buf]
# hack: the labelmap should be in the same folder with the tree file
# sometimes, the label in the tree file is different than the name in
# labelmap because there should be no whitespace in the label part. Thus,
# we have to use the labelmap to replace the labels in the tree file
true_labels = load_list_file(
op.join(op.dirname(label_tree_file), 'labelmap.txt'))
for i, true_label in enumerate(true_labels):
label_map[i][0] = true_label
if len(label_map[-1]) == 1:
label_map = label_map[:-1]
label_idx, label_parentidx = {}, {}
labels = []
idx = 0
for l_pi in label_map:
label, pi = l_pi[:2]
pi = int(pi)
label_parentidx[label] = pi
label_idx[label] = idx
idx = idx + 1
labels.append(label)
return label_idx, label_parentidx, labels
def load_labelmap(self):
return load_list_file(self.get_labelmap_file())
def load_noffsets(self):
logging.info('deprecated: pls generate it on the fly')
return load_list_file(self.get_noffsets_file())
def train(self):
torch.set_default_tensor_type('torch.cuda.FloatTensor')
kwargs = self.kwargs
from process_tsv import TSVDataset
tsv_dataset = TSVDataset(self.data)
tsv_file = tsv_dataset.get_data('train')
labelmap = tsv_dataset.get_labelmap_file()
cfg = kwargs
from qd_common import load_list_file
voc = {
'num_classes': len(load_list_file(labelmap)) + 1,
'lr_steps': (80000, 100000, 120000),
'max_iter': 120000,
'feature_maps': [38, 19, 10, 5, 3, 1],
'min_dim': 300,
'steps': [8, 16, 32, 64, 100, 300],
'min_sizes': [30, 60, 111, 162, 213, 264],
'max_sizes': [60, 111, 162, 213, 264, 315],
'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
'variance': [0.1, 0.2],
'clip': True,
'name': 'VOC',
}
for k in voc:
if k in cfg:
continue
cfg[k] = voc[k]
MEANS = (104, 117, 123)
dataset = TSVDetection(tsv_file=tsv_file,
labelmap=labelmap,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
net = ssd_net
if cfg['cuda']:
net = torch.nn.DataParallel(ssd_net)
cudnn.benchmark = True
vgg_weights = torch.load(cfg['save_folder'] + 'vgg16_reducedfc.pth')
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_weights)
if cfg['cuda']:
net = net.cuda()
if cfg['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=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3,
0.5, False, cfg['cuda'])
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
logging.info('Loading the dataset...')
epoch_size = len(dataset) // cfg['batch_size']
step_index = 0
data_loader = data.DataLoader(dataset,
cfg['batch_size'],
num_workers=cfg['num_workers'],
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
# create batch iterator
batch_iterator = iter(data_loader)
for iteration in range(cfg['start_iter'], cfg['max_iter']):
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, cfg['lr'], cfg['gamma'],
step_index)
t0 = time.time()
# load train data
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
if cfg['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]
data_loading_time = time.time() - t0
t0 = time.time()
# 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.data[0]
conf_loss += loss_c.data[0]
if iteration % 10 == 0:
logging.info('data loading time {}'.format(data_loading_time))
logging.info('timer: %.4f sec.' % (time.time() - t0))
logging.info('iter ' + repr(iteration) + ' || Loss: %.4f ||' %
(loss.data[0]))
if iteration != 0 and iteration % 500 == 0:
logging.info('Saving state, iter: {}'.format(iteration))
model_file = op.join(
self.output_folder, 'snapshot',
'model_iter_{}.pth.tar'.format(iteration + 1))
torch_save(ssd_net.state_dict(), model_file)
model_file = op.join(self.output_folder, 'snapshot',
'model_iter_{}.pth.tar'.format(iteration + 1))
torch_save(ssd_net.state_dict(), model_file)