def __init__(self, configer):
self.configer = configer
self.blob_helper = BlobHelper(configer)
self.seg_visualizer = SegVisualizer(configer)
self.seg_parser = SegParser(configer)
self.seg_model_manager = ModelManager(configer)
self.test_loader = TestDataLoader(configer)
self.device = torch.device(
'cpu' if self.configer.get('gpu') is None else 'cuda')
self.seg_net = None
self._init_model()
def build_global_network(self, learning_rate_input):
environment = Environment.create_environment(flags.env_type,
-1,
self.training_set,
shuffle=False)
self.global_network = ModelManager(-1, environment,
learning_rate_input, self.device)
# return gradient optimizer
return RMSPropApplier(learning_rate=learning_rate_input,
decay=flags.rmsp_alpha,
momentum=0.0,
epsilon=flags.rmsp_epsilon,
clip_norm=flags.grad_norm_clip,
device=self.device)
def __init__(self, configer):
self.configer = configer
self.batch_time = AverageMeter()
self.data_time = AverageMeter()
self.train_losses = DictAverageMeter()
self.val_losses = DictAverageMeter()
self.seg_running_score = SegRunningScore(configer)
self.seg_visualizer = SegVisualizer(configer)
self.seg_model_manager = ModelManager(configer)
self.seg_data_loader = DataLoader(configer)
self.seg_net = None
self.train_loader = None
self.val_loader = None
self.optimizer = None
self.scheduler = None
self.runner_state = dict()
self._init_model()
class FCNSegmentor(object):
"""
The class for Pose Estimation. Include train, val, val & predict.
"""
def __init__(self, configer):
self.configer = configer
self.batch_time = AverageMeter()
self.data_time = AverageMeter()
self.train_losses = DictAverageMeter()
self.val_losses = DictAverageMeter()
self.seg_running_score = SegRunningScore(configer)
self.seg_visualizer = SegVisualizer(configer)
self.seg_model_manager = ModelManager(configer)
self.seg_data_loader = DataLoader(configer)
self.seg_net = None
self.train_loader = None
self.val_loader = None
self.optimizer = None
self.scheduler = None
self.runner_state = dict()
self._init_model()
def _init_model(self):
self.seg_net = self.seg_model_manager.get_seg_model()
self.seg_net = RunnerHelper.load_net(self, self.seg_net)
self.optimizer, self.scheduler = Trainer.init(
self._get_parameters(), self.configer.get('solver'))
self.train_loader = self.seg_data_loader.get_trainloader()
self.val_loader = self.seg_data_loader.get_valloader()
self.loss = self.seg_model_manager.get_seg_loss()
def _get_parameters(self):
lr_1 = []
lr_10 = []
params_dict = dict(self.seg_net.named_parameters())
for key, value in params_dict.items():
if 'backbone' not in key:
lr_10.append(value)
else:
lr_1.append(value)
params = [{
'params': lr_1,
'lr': self.configer.get('solver', 'lr')['base_lr']
}, {
'params': lr_10,
'lr': self.configer.get('solver', 'lr')['base_lr'] * 1.0
}]
return params
def train(self):
"""
Train function of every epoch during train phase.
"""
self.seg_net.train()
start_time = time.time()
# Adjust the learning rate after every epoch.
for i, data_dict in enumerate(self.train_loader):
Trainer.update(self,
warm_list=(0, ),
solver_dict=self.configer.get('solver'))
self.data_time.update(time.time() - start_time)
# Forward pass.
data_dict = RunnerHelper.to_device(self, data_dict)
out = self.seg_net(data_dict)
# Compute the loss of the train batch & backward.
loss_dict = self.loss(out)
loss = loss_dict['loss']
self.train_losses.update(
{key: loss.item()
for key, loss in loss_dict.items()}, data_dict['img'].size(0))
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Update the vars of the train phase.
self.batch_time.update(time.time() - start_time)
start_time = time.time()
self.runner_state['iters'] += 1
# Print the log info & reset the states.
if self.runner_state['iters'] % self.configer.get(
'solver', 'display_iter') == 0:
Log.info(
'Train Epoch: {0}\tTrain Iteration: {1}\t'
'Time {batch_time.sum:.3f}s / {2}iters, ({batch_time.avg:.3f})\t'
'Data load {data_time.sum:.3f}s / {2}iters, ({data_time.avg:3f})\n'
'Learning rate = {4}\tLoss = {3}\n'.format(
self.runner_state['epoch'],
self.runner_state['iters'],
self.configer.get('solver', 'display_iter'),
self.train_losses.info(),
RunnerHelper.get_lr(self.optimizer),
batch_time=self.batch_time,
data_time=self.data_time))
self.batch_time.reset()
self.data_time.reset()
self.train_losses.reset()
if self.runner_state['iters'] % self.configer.get('solver.save_iters') == 0 \
and self.configer.get('local_rank') == 0:
RunnerHelper.save_net(self, self.seg_net)
if self.configer.get('solver', 'lr')['metric'] == 'iters' \
and self.runner_state['iters'] == self.configer.get('solver', 'max_iters'):
break
# Check to val the current model.
if self.runner_state['iters'] % self.configer.get('solver', 'test_interval') == 0 \
and not self.configer.get('network.distributed'):
self.val()
self.runner_state['epoch'] += 1
def val(self, data_loader=None):
"""
Validation function during the train phase.
"""
self.seg_net.eval()
start_time = time.time()
data_loader = self.val_loader if data_loader is None else data_loader
for j, data_dict in enumerate(data_loader):
data_dict = RunnerHelper.to_device(self, data_dict)
with torch.no_grad():
# Forward pass.
out = self.seg_net(data_dict)
loss_dict = self.loss(out)
# Compute the loss of the val batch.
out_dict, _ = RunnerHelper.gather(self, out)
self.val_losses.update(
{key: loss.item()
for key, loss in loss_dict.items()}, data_dict['img'].size(0))
self._update_running_score(out_dict['out'],
DCHelper.tolist(data_dict['meta']))
# Update the vars of the val phase.
self.batch_time.update(time.time() - start_time)
start_time = time.time()
self.runner_state['performance'] = self.seg_running_score.get_mean_iou(
)
self.runner_state['val_loss'] = self.val_losses.avg['loss']
RunnerHelper.save_net(
self,
self.seg_net,
performance=self.seg_running_score.get_mean_iou(),
val_loss=self.val_losses.avg['loss'])
# Print the log info & reset the states.
Log.info('Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t'
'Loss = {0}\n'.format(self.val_losses.info(),
batch_time=self.batch_time))
Log.info('Mean IOU: {}\n'.format(
self.seg_running_score.get_mean_iou()))
Log.info('Pixel ACC: {}\n'.format(
self.seg_running_score.get_pixel_acc()))
self.batch_time.reset()
self.val_losses.reset()
self.seg_running_score.reset()
self.seg_net.train()
def _update_running_score(self, pred, metas):
pred = pred.permute(0, 2, 3, 1)
for i in range(pred.size(0)):
border_size = metas[i]['border_wh']
ori_target = metas[i]['ori_target']
total_logits = cv2.resize(
pred[i, :border_size[1], :border_size[0]].cpu().numpy(),
tuple(metas[i]['ori_img_wh']),
interpolation=cv2.INTER_CUBIC)
labelmap = np.argmax(total_logits, axis=-1)
self.seg_running_score.update(labelmap[None], ori_target[None])
class Application(object):
def __init__(self):
# Training logger
self.training_logger = logging.getLogger('results')
if not os.path.isdir(flags.log_dir):
os.mkdir(flags.log_dir)
hdlr = logging.FileHandler(flags.log_dir + '/results.log')
formatter = logging.Formatter('%(asctime)s %(message)s')
hdlr.setFormatter(formatter)
self.training_logger.addHandler(hdlr)
self.training_logger.setLevel(logging.DEBUG)
# Test logger
self.test_logger = logging.getLogger('test')
if not os.path.isdir(flags.log_dir):
os.mkdir(flags.log_dir)
hdlr = logging.FileHandler(flags.log_dir + '/test.log')
# formatter = logging.Formatter('%(asctime)s %(message)s')
# hdlr.setFormatter(formatter)
self.test_logger.addHandler(hdlr)
self.test_logger.setLevel(logging.DEBUG)
# Build training and test set
self.training_set, self.test_set = self.get_set(
flags.preprocessed_dict + '.pkl')
# Shuffle training set (no need to shuffle test set)
np.random.shuffle(self.training_set)
# Initialize network
self.device = "/cpu:0"
if flags.use_gpu:
self.device = "/gpu:0"
config = tf.ConfigProto(log_device_placement=False,
allow_soft_placement=True) # prepare session
if flags.use_gpu:
config.gpu_options.allow_growth = True
self.sess = tf.Session(config=config)
self.global_t = 0
self.pause_requested = False
self.terminate_requested = False
self.build_network()
def get_set(self, path):
with open(path, 'rb') as f:
pkl = pickle.load(f)
return pkl['training_set'], pkl['test_set']
def build_network(self):
learning_rate_input = tf.placeholder("float")
grad_applier = self.build_global_network(learning_rate_input)
self.build_local_networks(learning_rate_input, grad_applier)
self.sess.run(tf.global_variables_initializer()
) # do it before loading checkpoint
self.load_checkpoint()
def build_global_network(self, learning_rate_input):
environment = Environment.create_environment(flags.env_type,
-1,
self.training_set,
shuffle=False)
self.global_network = ModelManager(-1, environment,
learning_rate_input, self.device)
# return gradient optimizer
return RMSPropApplier(learning_rate=learning_rate_input,
decay=flags.rmsp_alpha,
momentum=0.0,
epsilon=flags.rmsp_epsilon,
clip_norm=flags.grad_norm_clip,
device=self.device)
def build_local_networks(self, learning_rate_input, grad_applier):
initial_learning_rate = self.log_uniform(flags.initial_alpha_low,
flags.initial_alpha_high,
flags.initial_alpha_log_rate)
self.trainers = []
for i in range(flags.parallel_size):
trainer = Worker(i, self.sess, self.training_set,
self.global_network, self.device,
initial_learning_rate, learning_rate_input,
grad_applier)
self.trainers.append(trainer)
def log_uniform(self, lo, hi, rate):
log_lo = math.log(lo)
log_hi = math.log(hi)
v = log_lo * (1 - rate) + log_hi * rate
return math.exp(v)
def train_function(self, parallel_index, reset):
""" Train each environment. """
trainer = self.trainers[parallel_index]
# set start_time
trainer.set_start_time(self.start_time, reset)
print('Thread {0} started'.format(parallel_index))
while True:
if self.pause_requested:
break
if parallel_index == len(self.trainers) - 1:
if self.global_t > flags.max_time_step:
self.terminate_requested = True
if self.global_t > self.next_save_steps or self.terminate_requested:
self.save() # Save checkpoint
if self.terminate_requested or self.global_t > flags.max_time_step:
trainer.stop()
break
diff_global_t = trainer.process(step=self.global_t)
self.global_t += diff_global_t
# print global statistics
if trainer.terminal:
info = {}
for t in self.trainers:
for key in t.stats:
if not info.get(key):
info[key] = 0
info[key] += t.stats[key]
self.training_logger.info(
str([
key + "=" + str(value / len(self.trainers))
for key, value in sorted(info.items(),
key=lambda t: t[0])
])) # Print statistics
def test_function(self, parallel_index, tester):
lines = []
for id in range(parallel_index, len(self.test_set),
flags.parallel_size):
tester.prepare(id)
while not tester.terminal:
tester.process()
environment = tester.environment
annotation = copy.deepcopy(environment.annotation)
for i in range(len(environment.sentidoc)):
key = environment.get_task_by_index(i)
annotation[key] = environment.sentidoc[i]
line = '{0},{1},{2},{3},{4},{5},{6},{7}\n'.format(
annotation["id"], annotation["subjective"], annotation["opos"],
annotation["oneg"], annotation["ironic"], annotation["lpos"],
annotation["lneg"], annotation["topic"])
lines.append(line)
return lines
def test(self):
result_file = flags.log_dir + '/evaluation_' + str(
self.global_t) + '.csv'
if os.path.exists(result_file):
print('Test results already produced and evaluated for ' +
result_file)
return
print('Start testing')
threads = []
result_queue = Queue()
for i in range(flags.parallel_size): # parallel testing
tester = Worker(-1 - i,
self.sess,
self.test_set,
self.global_network,
self.device,
train=False)
thread = threading.Thread(target=lambda q, arg1, arg2: q.put(
self.test_function(arg1, arg2)),
args=(result_queue, i, tester))
thread.start()
threads.append(thread)
time.sleep(5)
for thread in threads: # wait for all threads to end
thread.join()
with open(result_file, "w",
encoding="utf-8") as file: # write results to file
while not result_queue.empty():
result = result_queue.get()
for line in result:
file.write(line)
print('End testing')
print('Test results saved in ' + result_file)
return self.evaluate(result_file)
def train(self):
# run training threads
self.train_threads = []
for i in range(flags.parallel_size):
self.train_threads.append(
threading.Thread(target=self.train_function, args=(i, True)))
signal.signal(signal.SIGINT, self.signal_handler)
# set start time
self.start_time = time.time() - self.wall_t
for t in self.train_threads:
t.start()
time.sleep(5)
print('Press Ctrl+C to stop')
signal.pause()
def load_checkpoint(self):
# init or load checkpoint with saver
self.saver = tf.train.Saver(var_list=self.global_network.get_vars(),
max_to_keep=0) # keep all checkpoints
checkpoint = tf.train.get_checkpoint_state(flags.checkpoint_dir)
if checkpoint and checkpoint.model_checkpoint_path:
self.saver.restore(self.sess, checkpoint.model_checkpoint_path)
print("checkpoint loaded:", checkpoint.model_checkpoint_path)
tokens = checkpoint.model_checkpoint_path.split("-")
# set global step
self.global_t = int(tokens[1])
print(">>> global step set: ", self.global_t)
# set wall time
wall_t_fname = flags.checkpoint_dir + '/' + 'wall_t.' + str(
self.global_t)
with open(wall_t_fname, 'r') as f:
self.wall_t = float(f.read())
self.next_save_steps = (
self.global_t + flags.save_interval_step
) // flags.save_interval_step * flags.save_interval_step
else:
print("Could not find old checkpoint")
# set wall time
self.wall_t = 0.0
self.next_save_steps = flags.save_interval_step
def save(self):
""" Save checkpoint.
Called from thread-0.
"""
self.pause_requested = True
for (i, t) in enumerate(
self.train_threads): # Wait for all other threads to stop
if i != len(self.train_threads) - 1: # cannot join current thread
t.join()
# Save
if not os.path.exists(flags.checkpoint_dir):
os.mkdir(flags.checkpoint_dir)
# Write wall time
wall_t = time.time() - self.start_time
wall_t_fname = flags.checkpoint_dir + '/wall_t.' + str(self.global_t)
with open(wall_t_fname, 'w') as f:
f.write(str(wall_t))
print('Start saving')
self.saver.save(self.sess,
flags.checkpoint_dir + '/checkpoint',
global_step=self.global_t)
print('End saving')
# Test
test_result = self.test()
# Restart training
if not self.terminate_requested:
self.pause_requested = False
self.next_save_steps += flags.save_interval_step
# Restart other threads
for i in range(flags.parallel_size):
if i != len(self.train_threads
) - 1: # current thread is already running
thread = threading.Thread(target=self.train_function,
args=(i, False))
self.train_threads[i] = thread
thread.start()
def evaluate(self, result_file):
print('Start evaluating')
verbose = True
self.test_logger.info(datetime.now())
# read gold standard and populate the count matrix
gold = dict()
gold_counts = {
'subj': {
'0': 0,
'1': 0
},
'opos': {
'0': 0,
'1': 0
},
'oneg': {
'0': 0,
'1': 0
},
'iro': {
'0': 0,
'1': 0
},
'lpos': {
'0': 0,
'1': 0
},
'lneg': {
'0': 0,
'1': 0
}
}
with open(flags.test_set_path) as f:
for line in f:
raw = line.rstrip().split(',')
id = str(raw[0].replace('"', ''))
subj = str(raw[1].replace('"', ''))
opos = str(raw[2].replace('"', ''))
oneg = str(raw[3].replace('"', ''))
iro = str(raw[4].replace('"', ''))
lpos = str(raw[5].replace('"', ''))
lneg = str(raw[6].replace('"', ''))
top = str(raw[7].replace('"', ''))
#id, subj, opos, oneg, iro, lpos, lneg, top = map(lambda x: x[1:-1], line.rstrip().split(','))
gold[id] = {
'subj': subj,
'opos': opos,
'oneg': oneg,
'iro': iro,
'lpos': lpos,
'lneg': lneg
}
gold_counts['subj'][subj] += 1
gold_counts['opos'][opos] += 1
gold_counts['oneg'][oneg] += 1
gold_counts['iro'][iro] += 1
gold_counts['lpos'][lpos] += 1
gold_counts['lneg'][lneg] += 1
# read result data
result = dict()
with open(result_file) as f:
for line in f:
raw = line.rstrip().split(',')
id = str(raw[0].replace('"', ''))
subj = str(raw[1].replace('"', ''))
opos = str(raw[2].replace('"', ''))
oneg = str(raw[3].replace('"', ''))
iro = str(raw[4].replace('"', ''))
lpos = str(raw[5].replace('"', ''))
lneg = str(raw[6].replace('"', ''))
top = str(raw[7].replace('"', ''))
result[id] = {
'subj': subj,
'opos': opos,
'oneg': oneg,
'iro': iro
}
task_f1 = {}
# evaluation: single classes
for task in [
'subj', 'opos', 'oneg', 'iro'
]: #add 'lpos' and 'lneg' if you want to measure literal polairty
# table header
if verbose: self.test_logger.info("\ntask: {}".format(task))
if verbose:
self.test_logger.info(
"prec. 0\trec. 0\tF-sc. 0\tprec. 1\trec. 1\tF-sc. 1\tF-sc."
)
correct = {'0': 0, '1': 0}
assigned = {'0': 0, '1': 0}
precision = {'0': 0.0, '1': 0.0}
recall = {'0': 0.0, '1': 0.0}
fscore = {'0': 0.0, '1': 0.0}
# count the labels
for id, gold_labels in gold.items():
if (not id in result) or result[id][task] == '':
pass
else:
assigned[result[id][task]] += 1
if gold_labels[task] == result[id][task]:
correct[result[id][task]] += 1
# compute precision, recall and F-score
for label in ['0', '1']:
try:
precision[label] = float(correct[label]) / float(
assigned[label])
recall[label] = float(correct[label]) / float(
gold_counts[task][label])
fscore[label] = (2.0 * precision[label] * recall[label]
) / (precision[label] + recall[label])
except:
# if a team doesn't participate in a task it gets default 0 F-score
fscore[label] = 0.0
task_f1[task] = (fscore['0'] + fscore['1']) / 2.0
# write down the table
self.test_logger.info(
"{0:.4f}\t{1:.4f}\t{2:.4f}\t{3:.4f}\t{4:.4f}\t{5:.4f}\t{6:.4f}"
.format(precision['0'], recall['0'], fscore['0'],
precision['1'], recall['1'], fscore['1'],
task_f1[task]))
# polarity evaluation needs a further step
if verbose: self.test_logger.info("\ntask: polarity")
if verbose: self.test_logger.info("Combined F-score")
correct = {'opos': {'0': 0, '1': 0}, 'oneg': {'0': 0, '1': 0}}
assigned = {'opos': {'0': 0, '1': 0}, 'oneg': {'0': 0, '1': 0}}
precision = {
'opos': {
'0': 0.0,
'1': 0.0
},
'oneg': {
'0': 0.0,
'1': 0.0
}
}
recall = {'opos': {'0': 0.0, '1': 0.0}, 'oneg': {'0': 0.0, '1': 0.0}}
fscore = {'opos': {'0': 0.0, '1': 0.0}, 'oneg': {'0': 0.0, '1': 0.0}}
# count the labels
for id, gold_labels in gold.items():
for cl in ['opos', 'oneg']:
if (not id in result) or result[id][cl] == '':
pass
else:
assigned[cl][result[id][cl]] += 1
if gold_labels[cl] == result[id][cl]:
correct[cl][result[id][cl]] += 1
# compute precision, recall and F-score
for cl in ['opos', 'oneg']:
for label in ['0', '1']:
try:
precision[cl][label] = float(correct[cl][label]) / float(
assigned[cl][label])
recall[cl][label] = float(correct[cl][label]) / float(
gold_counts[cl][label])
fscore[cl][label] = float(
2.0 * precision[cl][label] *
recall[cl][label]) / float(precision[cl][label] +
recall[cl][label])
except:
fscore[cl][label] = 0.0
fscore_pos = (fscore['opos']['0'] + fscore['opos']['1']) / 2.0
fscore_neg = (fscore['oneg']['0'] + fscore['oneg']['1']) / 2.0
# write down the table
task_f1["polarity"] = (fscore_pos + fscore_neg) / 2.0
self.test_logger.info("{0:.4f}".format(task_f1["polarity"]))
print('End evaluating')
return task_f1
def signal_handler(self, signal, frame):
print('You pressed Ctrl+C!')
self.terminate_requested = True
def find_best_parameters(self, save_path, save_graph=False, verbose=True):
best_accuracy = -1
best_model = None
model_manager = ModelManager(None, verbose=verbose)
for sample in range(self.num_samples):
self.sample_parameters()
if verbose:
print('Evaluating model:\n{}'.format(self.model_name))
try:
model_params = {
'train_file': self.train_file,
'validation_file': self.validation_file,
'test_file': self.test_file,
'saved_model_folder': self.saved_model_folder,
'num_train': self.num_train,
'num_validation': self.num_validation,
'num_test': self.num_test,
'use_validation': self.use_validation,
'use_mc_dropout': self.use_mc_dropout,
'graphs_dir': self.graphs_dir,
'model_name': self.model_name,
'tensorboard_dir': self.tensorboard_dir,
'embedding_file': self.embedding_file,
'embedding_pickle': self.embedding_pickle,
'learning_rate': self.learning_rate,
'batch_size': self.batch_size,
'num_epochs': self.num_epochs,
'perform_shuffle': self.perform_shuffle,
'embed_size': self.embed_size,
'num_units': self.num_units,
'num_classes': self.num_classes,
'recurrent_output_dropout': self.recurrent_output_dropout,
'recurrent_state_dropout': self.recurrent_state_dropout,
'embedding_dropout': self.embedding_dropout,
'clip_gradients': self.clip_gradients,
'max_norm': self.max_norm,
'weight_decay': self.weight_decay,
'bucket_width': self.bucket_width,
'num_buckets': self.num_buckets,
'use_test': self.use_test,
'save_graph': self.save_graph,
'should_save': False
}
model_manager.model_params = model_params
accuracy, _, _, _ = model_manager.run_model()
model_manager.reset_graph()
if accuracy > best_accuracy:
best_accuracy = accuracy
best_model = self.model_name
except tf.errors.InvalidArgumentError as e:
print('Error running model: {}'.format(str(e)))
model_manager.reset_graph()
continue
except tf.errors.ResourceExhaustedError as e:
print('Error running model: {}'.format(str(e)))
model_manager.reset_graph()
continue
if verbose:
print('Best model:\n{}'.format(best_model))
print('Accuracy on validation set: {}'.format(best_accuracy))
if verbose:
print('Saving best model parameters ...')
if best_model:
self.save_parameters(best_model, best_accuracy, save_path)
class FCNSegmentorTest(object):
def __init__(self, configer):
self.configer = configer
self.blob_helper = BlobHelper(configer)
self.seg_visualizer = SegVisualizer(configer)
self.seg_parser = SegParser(configer)
self.seg_model_manager = ModelManager(configer)
self.test_loader = TestDataLoader(configer)
self.device = torch.device(
'cpu' if self.configer.get('gpu') is None else 'cuda')
self.seg_net = None
self._init_model()
def _init_model(self):
self.seg_net = self.seg_model_manager.get_seg_model()
self.seg_net = RunnerHelper.load_net(self, self.seg_net)
self.seg_net.eval()
def test(self, test_dir, out_dir):
for _, data_dict in enumerate(
self.test_loader.get_testloader(test_dir=test_dir)):
total_logits = None
if self.configer.get('test', 'mode') == 'ss_test':
total_logits = self.ss_test(data_dict)
elif self.configer.get('test', 'mode') == 'sscrop_test':
total_logits = self.sscrop_test(data_dict,
params_dict=self.configer.get(
'test', 'sscrop_test'))
elif self.configer.get('test', 'mode') == 'ms_test':
total_logits = self.ms_test(data_dict,
params_dict=self.configer.get(
'test', 'ms_test'))
elif self.configer.get('test', 'mode') == 'mscrop_test':
total_logits = self.mscrop_test(data_dict,
params_dict=self.configer.get(
'test', 'mscrop_test'))
else:
Log.error('Invalid test mode:{}'.format(
self.configer.get('test', 'mode')))
exit(1)
meta_list = DCHelper.tolist(data_dict['meta'])
for i in range(len(meta_list)):
label_map = np.argmax(total_logits[i], axis=-1)
label_img = np.array(label_map, dtype=np.uint8)
ori_img_bgr = ImageHelper.read_image(meta_list[i]['img_path'],
tool='cv2',
mode='BGR')
image_canvas = self.seg_parser.colorize(
label_img, image_canvas=ori_img_bgr)
ImageHelper.save(image_canvas,
save_path=os.path.join(
out_dir, 'vis/{}.png'.format(
meta_list[i]['filename'])))
if self.configer.get('data.label_list',
default=None) is not None:
label_img = self.__relabel(label_img)
if self.configer.get('data.reduce_zero_label', default=False):
label_img = label_img + 1
label_img = label_img.astype(np.uint8)
label_img = Image.fromarray(label_img, 'P')
label_path = os.path.join(
out_dir, 'label/{}.png'.format(meta_list[i]['filename']))
Log.info('Label Path: {}'.format(label_path))
ImageHelper.save(label_img, label_path)
def ss_test(self, in_data_dict):
data_dict = self.blob_helper.get_blob(in_data_dict, scale=1.0)
results = self._predict(data_dict)
return results
def ms_test(self, in_data_dict, params_dict):
total_logits = [
np.zeros((meta['ori_img_size'][1], meta['ori_img_size'][0],
self.configer.get('data', 'num_classes')), np.float32)
for meta in DCHelper.tolist(in_data_dict['meta'])
]
for scale in params_dict['scale_search']:
data_dict = self.blob_helper.get_blob(in_data_dict, scale=scale)
results = self._predict(data_dict)
for i in range(len(total_logits)):
total_logits[i] += results[i]
for scale in params_dict['scale_search']:
data_dict = self.blob_helper.get_blob(in_data_dict,
scale=scale,
flip=True)
results = self._predict(data_dict)
for i in range(len(total_logits)):
total_logits[i] += results[i][:, ::-1]
return total_logits
def sscrop_test(self, in_data_dict, params_dict):
data_dict = self.blob_helper.get_blob(in_data_dict, scale=1.0)
if any(image.size()[2] < params_dict['crop_size'][0]
or image.size()[1] < params_dict['crop_size'][1]
for image in DCHelper.tolist(data_dict['img'])):
results = self._predict(data_dict)
else:
results = self._crop_predict(data_dict, params_dict['crop_size'],
params_dict['crop_stride_ratio'])
return results
def mscrop_test(self, in_data_dict, params_dict):
total_logits = [
np.zeros((meta['ori_img_size'][1], meta['ori_img_size'][0],
self.configer.get('data', 'num_classes')), np.float32)
for meta in DCHelper.tolist(in_data_dict['meta'])
]
for scale in params_dict['scale_search']:
data_dict = self.blob_helper.get_blob(in_data_dict, scale=scale)
if any(image.size()[2] < params_dict['crop_size'][0]
or image.size()[1] < params_dict['crop_size'][1]
for image in DCHelper.tolist(data_dict['img'])):
results = self._predict(data_dict)
else:
results = self._crop_predict(data_dict,
params_dict['crop_size'],
params_dict['crop_stride_ratio'])
for i in range(len(total_logits)):
total_logits[i] += results[i]
for scale in params_dict['scale_search']:
data_dict = self.blob_helper.get_blob(in_data_dict,
scale=scale,
flip=True)
if any(image.size()[2] < params_dict['crop_size'][0]
or image.size()[1] < params_dict['crop_size'][1]
for image in DCHelper.tolist(data_dict['img'])):
results = self._predict(data_dict)
else:
results = self._crop_predict(data_dict,
params_dict['crop_size'],
params_dict['crop_stride_ratio'])
for i in range(len(total_logits)):
total_logits[i] += results[i][:, ::-1]
return total_logits
def _crop_predict(self, data_dict, crop_size, crop_stride_ratio):
split_batch = list()
height_starts_list = list()
width_starts_list = list()
hw_list = list()
for image in DCHelper.tolist(data_dict['img']):
height, width = image.size()[1:]
hw_list.append([height, width])
np_image = image.squeeze(0).permute(1, 2, 0).cpu().numpy()
height_starts = self._decide_intersection(height, crop_size[1],
crop_stride_ratio)
width_starts = self._decide_intersection(width, crop_size[0],
crop_stride_ratio)
split_crops = []
for height in height_starts:
for width in width_starts:
image_crop = np_image[height:height + crop_size[1],
width:width + crop_size[0]]
split_crops.append(image_crop[np.newaxis, :])
height_starts_list.append(height_starts)
width_starts_list.append(width_starts)
split_crops = np.concatenate(
split_crops,
axis=0) # (n, crop_image_size, crop_image_size, 3)
inputs = torch.from_numpy(split_crops).permute(0, 3, 1,
2).to(self.device)
split_batch.append(inputs)
assert len(split_batch) == torch.cuda.device_count(
), 'Only support one image per gpu.'
out_list = list()
with torch.no_grad():
results = self.seg_net(
dict(img=DCHelper.todc(split_batch,
stack=False,
samples_per_gpu=True,
concat=True)))
results = results if isinstance(results,
(list, tuple)) else [results]
for res in results:
out_list.append(res['out'].permute(0, 2, 3, 1).cpu().numpy())
total_logits = [
np.zeros((hw[0], hw[1], self.configer.get('data', 'num_classes')),
np.float32) for hw in hw_list
]
count_predictions = [
np.zeros((hw[0], hw[1], self.configer.get('data', 'num_classes')),
np.float32) for hw in hw_list
]
for i in range(len(height_starts_list)):
index = 0
for height in height_starts_list[i]:
for width in width_starts_list[i]:
total_logits[i][height:height + crop_size[1], width:width +
crop_size[0]] += out_list[i][index]
count_predictions[i][height:height + crop_size[1],
width:width + crop_size[0]] += 1
index += 1
for i in range(len(total_logits)):
total_logits[i] /= count_predictions[i]
for i, meta in enumerate(DCHelper.tolist(data_dict['meta'])):
total_logits[i] = cv2.resize(
total_logits[i][:meta['border_wh'][1], :meta['border_wh'][0]],
tuple(meta['ori_img_size']),
interpolation=cv2.INTER_CUBIC)
return total_logits
def _decide_intersection(self, total_length, crop_length,
crop_stride_ratio):
stride = int(crop_length *
crop_stride_ratio) # set the stride as the paper do
times = (total_length - crop_length) // stride + 1
cropped_starting = []
for i in range(times):
cropped_starting.append(stride * i)
if total_length - cropped_starting[-1] > crop_length:
cropped_starting.append(total_length -
crop_length) # must cover the total image
return cropped_starting
def _predict(self, data_dict):
with torch.no_grad():
total_logits = list()
results = self.seg_net(data_dict)
results = results if isinstance(results,
(list, tuple)) else [results]
for res in results:
assert res['out'].size(
0) == 1, 'Only support one image per gpu.'
total_logits.append(res['out'].squeeze(0).permute(
1, 2, 0).cpu().numpy())
for i, meta in enumerate(DCHelper.tolist(data_dict['meta'])):
total_logits[i] = cv2.resize(
total_logits[i]
[:meta['border_wh'][1], :meta['border_wh'][0]],
tuple(meta['ori_img_size']),
interpolation=cv2.INTER_CUBIC)
return total_logits
def __relabel(self, label_map):
height, width = label_map.shape
label_dst = np.zeros((height, width), dtype=np.uint8)
for i in range(self.configer.get('data', 'num_classes')):
label_dst[label_map == i] = self.configer.get(
'data', 'label_list')[i]
label_dst = np.array(label_dst, dtype=np.uint8)
return label_dst
def __init__(self,
thread_index,
session,
annotated_wordbag,
global_network,
device,
initial_learning_rate=None,
learning_rate_input=None,
grad_applier=None,
train=True):
self.train = train
self.thread_index = thread_index
self.sess = session
self.global_network = global_network
self.environment = Environment.create_environment(flags.env_type,
self.thread_index,
annotated_wordbag,
shuffle=self.train)
self.device = device
# build network
if self.train:
self.local_network = ModelManager(self.thread_index,
self.environment,
learning_rate_input, self.device)
self.apply_gradients, self.sync = self.local_network.initialize(
self.global_network, grad_applier)
self.initial_learning_rate = initial_learning_rate
else:
self.local_network = self.global_network
self.terminal = True
self.local_t = 0
self.prev_local_t = 0
#logs
if self.train:
# main log directory
if not os.path.exists(flags.log_dir):
os.makedirs(flags.log_dir)
# episode result
self.result_log_dir = flags.log_dir + "/thread" + str(
self.thread_index)
if not os.path.exists(self.result_log_dir):
os.makedirs(self.result_log_dir)
# perfomance
if not os.path.exists(flags.log_dir + "/performance"):
os.makedirs(flags.log_dir + "/performance")
formatter = logging.Formatter('%(asctime)s %(message)s')
# info logger
self.info_logger = logging.getLogger('info_' + str(thread_index))
hdlr = logging.FileHandler(flags.log_dir + '/performance/info_' +
str(thread_index) + '.log')
hdlr.setFormatter(formatter)
self.info_logger.addHandler(hdlr)
self.info_logger.setLevel(logging.DEBUG)
# reward logger
self.reward_logger = logging.getLogger('reward_' +
str(thread_index))
hdlr = logging.FileHandler(flags.log_dir + '/performance/reward_' +
str(thread_index) + '.log')
hdlr.setFormatter(formatter)
self.reward_logger.addHandler(hdlr)
self.reward_logger.setLevel(logging.DEBUG)
self.max_reward = float("-inf")
self.update_statistics()
class Worker(object):
def __init__(self,
thread_index,
session,
annotated_wordbag,
global_network,
device,
initial_learning_rate=None,
learning_rate_input=None,
grad_applier=None,
train=True):
self.train = train
self.thread_index = thread_index
self.sess = session
self.global_network = global_network
self.environment = Environment.create_environment(flags.env_type,
self.thread_index,
annotated_wordbag,
shuffle=self.train)
self.device = device
# build network
if self.train:
self.local_network = ModelManager(self.thread_index,
self.environment,
learning_rate_input, self.device)
self.apply_gradients, self.sync = self.local_network.initialize(
self.global_network, grad_applier)
self.initial_learning_rate = initial_learning_rate
else:
self.local_network = self.global_network
self.terminal = True
self.local_t = 0
self.prev_local_t = 0
#logs
if self.train:
# main log directory
if not os.path.exists(flags.log_dir):
os.makedirs(flags.log_dir)
# episode result
self.result_log_dir = flags.log_dir + "/thread" + str(
self.thread_index)
if not os.path.exists(self.result_log_dir):
os.makedirs(self.result_log_dir)
# perfomance
if not os.path.exists(flags.log_dir + "/performance"):
os.makedirs(flags.log_dir + "/performance")
formatter = logging.Formatter('%(asctime)s %(message)s')
# info logger
self.info_logger = logging.getLogger('info_' + str(thread_index))
hdlr = logging.FileHandler(flags.log_dir + '/performance/info_' +
str(thread_index) + '.log')
hdlr.setFormatter(formatter)
self.info_logger.addHandler(hdlr)
self.info_logger.setLevel(logging.DEBUG)
# reward logger
self.reward_logger = logging.getLogger('reward_' +
str(thread_index))
hdlr = logging.FileHandler(flags.log_dir + '/performance/reward_' +
str(thread_index) + '.log')
hdlr.setFormatter(formatter)
self.reward_logger.addHandler(hdlr)
self.reward_logger.setLevel(logging.DEBUG)
self.max_reward = float("-inf")
self.update_statistics()
def update_statistics(self):
self.stats = self.environment.get_statistics()
self.stats.update(self.local_network.get_statistics())
def prepare(self, episode_id=None): # initialize a new episode
self.terminal = False
self.environment.reset(episode_id)
self.local_network.reset()
def stop(self): # stop current episode
self.environment.stop()
def _anneal_learning_rate(self, global_step): # anneal learning rate
learning_rate = self.initial_learning_rate * (
flags.max_time_step - global_step) / flags.max_time_step
if learning_rate < 0.0:
learning_rate = 0.0
return learning_rate
def set_start_time(self, start_time, reset):
self.start_time = start_time
if reset:
self.local_network.init_train_count()
def _print_log(self, step):
# if self.local_t - self.prev_local_t >= PERFORMANCE_LOG_INTERVAL):
# self.prev_local_t += PERFORMANCE_LOG_INTERVAL
# elapsed_time = time.time() - self.start_time
# steps_per_sec = step / elapsed_time
# print("### Performance : {} STEPS in {:.0f} sec. {:.0f} STEPS/sec. {:.2f}M STEPS/hour".format( step, elapsed_time, steps_per_sec, steps_per_sec * 3600 / 1000000.))
# Print statistics
self.reward_logger.info(
str([
"{0}={1}".format(key, value)
for key, value in self.stats.items()
]))
# Print match results
print_result = False
if flags.show_all_screenshots:
print_result = True
elif flags.show_best_screenshots:
if self.environment.episode_reward > self.max_reward:
self.max_reward = self.environment.episode_reward
print_result = True
if print_result: # show episodes insides
file = open(
self.result_log_dir +
'/reward({0})_epoch({1})_step({2}).log'.format(
self.environment.episode_reward, self.environment.epoch,
step), "w")
file.write('Annotation: {0}\n'.format([
"{0}={1}".format(key, value)
for key, value in sorted(self.environment.annotation.items(),
key=lambda t: t[0])
]))
file.write('Prediction: {0}\n'.format([
"{0}={1}".format(key, value) for key, value in sorted(
self.environment.get_labeled_prediction().items(),
key=lambda t: t[0])
]))
file.write('Reward: {0}\n'.format(
self.environment.compute_reward()))
file.close()
# run simulations and build a batch for training
def _build_batch(self, step):
batch = {}
if self.train: # init batch
batch["states"] = []
batch["actions"] = []
batch["concat"] = []
batch["cumulative_rewards"] = []
batch["advantages"] = []
batch["start_lstm_state"] = []
for i in range(self.local_network.model_size):
for key in batch:
batch[key].append(collections.deque())
batch["start_lstm_state"][i] = self.local_network.get_model(
i).lstm_state_out
agent_id_list = collections.deque()
agent_reward_list = collections.deque()
agent_value_list = collections.deque()
manager_value_list = collections.deque()
t = 0
while t < flags.local_t_max and not self.terminal:
t += 1
prediction = self.environment.sentidoc
state = self.environment.get_state()
agent_id, manager_policy, manager_value = self.local_network.get_agentID_by_state(
sess=self.sess, state=[state], concat=[prediction])
agent = self.local_network.get_model(agent_id)
agent_policy, agent_value = agent.run_policy_and_value(
sess=self.sess, state=[state], concat=[prediction])
action = self.environment.choose_action(agent_policy)
reward, self.terminal = self.environment.process_action(action)
self.local_t += 1
if self.train: # populate batch
if self.terminal: # update statistics
self.update_statistics(
) # required before assigning manager reward
# Populate agent batch
batch["states"][agent_id].append(state)
batch["actions"][agent_id].append(
agent.get_action_vector(action))
batch["concat"][agent_id].append(prediction)
agent_id_list.appendleft(agent_id) # insert on top
agent_reward_list.appendleft(reward) # insert on top
agent_value_list.appendleft(agent_value) # insert on top
# Populate manager batch
if self.local_network.has_manager:
batch["states"][0].append(state)
batch["actions"][0].append(
self.local_network.manager.get_action_vector(agent_id -
1))
batch["concat"][0].append(prediction)
manager_value_list.appendleft(
manager_value) # insert on top
if (self.local_t % LOG_INTERVAL == 0):
self.info_logger.info(
"actions={0} value={1} agent={2}".format(
agent_policy, agent_value, agent_id))
# build cumulative reward
if self.train:
cumulative_reward = 0.0
# if the episode has not terminated, bootstrap the value from the last state
if not self.terminal:
prediction = self.environment.sentidoc
state = self.environment.get_state()
agent_id, manager_policy, manager_value = self.local_network.get_agentID_by_state(
sess=self.sess, state=[state], concat=[prediction])
agent = self.local_network.get_model(agent_id)
agent_value = agent.run_value(self.sess,
state=[state],
concat=[prediction])
if self.local_network.has_manager:
cumulative_reward = manager_value if abs(
manager_value
) > abs(
agent_value
) else agent_value # should prevent value over-estimation
else:
cumulative_reward = agent_value
if self.local_network.has_manager:
for (agent_id, agent_reward, agent_value,
manager_value) in zip(agent_id_list, agent_reward_list,
agent_value_list,
manager_value_list):
value = manager_value if abs(manager_value) > abs(
agent_value) else agent_value
cumulative_reward = agent_reward + flags.gamma * cumulative_reward
batch["cumulative_rewards"][agent_id].appendleft(
cumulative_reward) # insert on top
batch["advantages"][agent_id].appendleft(
cumulative_reward - value) # insert on top
batch["cumulative_rewards"][0].appendleft(
cumulative_reward) # insert on top
batch["advantages"][0].appendleft(cumulative_reward -
value) # insert on top
else:
for (agent_id, agent_reward,
agent_value) in zip(agent_id_list, agent_reward_list,
agent_value_list):
cumulative_reward = agent_reward + flags.gamma * cumulative_reward
batch["cumulative_rewards"][agent_id].appendleft(
cumulative_reward) # insert on top
batch["advantages"][agent_id].appendleft(
cumulative_reward - agent_value) # insert on top
return batch
def process(self, step=0):
if self.terminal:
self.prepare()
start_local_t = self.local_t
# Copy weights from shared to local
if self.train:
learning_rate = []
agents_count = self.local_network.model_size - 1
for i in range(self.local_network.model_size):
self.sess.run(self.sync[i])
rate = self._anneal_learning_rate(
self.local_network.get_model(i).train_count)
# manager learning rate should be the highest
if i > 0:
rate /= agents_count
learning_rate.append(rate)
# Build feed dictionary
batch_dict = self._build_batch(step)
# Pupulate the feed dictionary
if self.train:
for i in range(self.local_network.model_size):
if len(batch_dict["states"][i]) > 0:
agent = self.local_network.get_model(i)
agent.train(self.sess, self.apply_gradients[i],
learning_rate[i], batch_dict["states"][i],
batch_dict["actions"][i],
batch_dict["advantages"][i],
batch_dict["cumulative_rewards"][i],
batch_dict["start_lstm_state"][i],
batch_dict["concat"][i])
self._print_log(step)
diff_local_t = self.local_t - start_local_t
return diff_local_t # local steps amount
default=False,
help=
'Classify or train on greyscale images, if False then use RGB color scheme'
)
parser.add_argument(
'-m',
'--model-name',
nargs='?',
type=str,
default='base',
help=
'Provide the name of trained model. Possible values: base, lenet, stridenet'
)
args = vars(parser.parse_args())
mm = ModelManager(args=args)
mm.get_model()
if args['directory'] is not None:
validation_dir_path = Path(args['directory'])
if validation_dir_path.is_dir():
files = glob(os.path.join(args['directory'], '*.jpg'))
imgs = []
for f in files:
if args['greyscale']:
img = CoinImage(
img_arr=cv2.imread(f, cv2.IMREAD_GRAYSCALE))
else:
img = CoinImage(img_arr=cv2.imread(f, cv2.IMREAD_COLOR))
imgs.append(img)
mm.classify_coin_images(imgs)
else: