def __init__(self, model_dir=None):
self.tp_num = 0
self.pred_num = 0
self.gold_num = 0
if model_dir:
self.model_dir = os.path.join(Config.train.model_dir,
'eval_' + model_dir)
else:
self.model_dir = os.path.join(Config.train.model_dir, 'eval')
self.label_dict = data_loader.load_label()
def main():
## load data
# data_arr_01 = data_loader.load_data(r'toydata/data.txt')
# label_vec = data_loader.load_label(r'toydata/label.txt')
data_arr_01 = data_loader.load_data(r'data/uestc_pgb/SF01/vib_data_1.txt')
# data_arr_03 = data_loader.load_data('data/pgb/SF03/vib_data_1.txt')
# data_arr_01 = data_loader.resample_arr(data_arr_01, num=240) # add for Ince's model
# data_arr_03 = data_loader.resample_arr(data_arr_03, num=240) # add for Ince's model
# data_arr_01, _ = data_loader.fft_arr(data_arr_01) # add for fft wdcnn
# data_arr_03, _ = data_loader.fft_arr(data_arr_03) # add for fft wdcnn
# data_arr_01 = data_loader.stft_arr(data_arr_01) # add for stft-LeNet
# data_arr_03 = data_loader.stft_arr(data_arr_03)
label_vec = data_loader.load_label(r'data/uestc_pgb/SF01/label_vec.txt')
trainset_01, testset_01 = data_loader.split_set(data_arr_01, label_vec)
# trainset_03, testset_03 = data_loader.split_set(data_arr_03, label_vec)
train_loader = data_utils.DataLoader(dataset=trainset_01,
batch_size=512,
shuffle=True,
num_workers=2)
test_loader = data_utils.DataLoader(dataset=testset_01,
batch_size=512,
shuffle=True,
num_workers=2)
print('Number of training samples: {}'.format(len(train_loader.dataset)))
print('Number of testing samples: {}'.format(len(test_loader.dataset)))
print()
## make models
model = wdcnn.Net(1, 5)
## train
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), weight_decay=0.0001)
best_model, loss_curve = iter_utils.train(model,
train_loader,
criterion,
optimizer,
init_lr=0.0001,
decay_epoch=5,
n_epoch=10,
use_cuda=False)
# test
test_accuracy = iter_utils.test(best_model, test_loader)
print('Test accuracy: {:.4f}%'.format(100 * test_accuracy))
def train_model(data_folder,
data_name,
level,
model_name,
is_aspect_term=True):
config.data_folder = data_folder
config.data_name = data_name
if not os.path.exists(os.path.join(config.checkpoint_dir, data_folder)):
os.makedirs(os.path.join(config.checkpoint_dir, data_folder))
config.level = level
config.model_name = model_name
config.is_aspect_term = is_aspect_term
config.init_input()
config.exp_name = '{}_{}_wv_{}'.format(model_name, level,
config.word_embed_type)
config.exp_name = config.exp_name + '_update' if config.word_embed_trainable else config.exp_name + '_fix'
if config.use_aspect_input:
config.exp_name += '_aspv_{}'.format(config.aspect_embed_type)
config.exp_name = config.exp_name + '_update' if config.aspect_embed_trainable else config.exp_name + '_fix'
if config.use_elmo:
config.exp_name += '_elmo_alone_{}_mode_{}_{}'.format(
config.use_elmo_alone, config.elmo_output_mode,
'update' if config.elmo_trainable else 'fix')
print(config.exp_name)
model = SentimentModel(config)
test_input = load_input_data(
data_folder, 'test', level, config.use_text_input,
config.use_text_input_l, config.use_text_input_r,
config.use_text_input_r_with_pad, config.use_aspect_input,
config.use_aspect_text_input, config.use_loc_input,
config.use_offset_input, config.use_mask)
test_label = load_label(data_folder, 'test')
if not os.path.exists(
os.path.join(config.checkpoint_dir, '%s/%s.hdf5' %
(data_folder, config.exp_name))):
start_time = time.time()
train_input = load_input_data(
data_folder, 'train', level, config.use_text_input,
config.use_text_input_l, config.use_text_input_r,
config.use_text_input_r_with_pad, config.use_aspect_input,
config.use_aspect_text_input, config.use_loc_input,
config.use_offset_input, config.use_mask)
train_label = load_label(data_folder, 'train')
valid_input = load_input_data(
data_folder, 'valid', level, config.use_text_input,
config.use_text_input_l, config.use_text_input_r,
config.use_text_input_r_with_pad, config.use_aspect_input,
config.use_aspect_text_input, config.use_loc_input,
config.use_offset_input, config.use_mask)
valid_label = load_label(data_folder, 'valid')
train_combine_valid_input = []
for i in range(len(train_input)):
train_combine_valid_input.append(train_input[i] + valid_input[i])
train_combine_valid_label = train_label + valid_label
model.train(train_combine_valid_input, train_combine_valid_label,
test_input, test_label)
elapsed_time = time.time() - start_time
print('training time:',
time.strftime("%H:%M:%S", time.gmtime(elapsed_time)))
# load the best model
model.load()
# print('score over valid data...')
# model.score(valid_input, valid_label)
print('score over test data...')
model.score(test_input, test_label)
def segment_main(**kwargs):
mode = kwargs['mode']
mu = FLAGS.mean
if FLAGS.ignore_class_bg:
nSketchClasses = FLAGS.nSketchClasses - 1
print('Ignore BG;', nSketchClasses, 'classes')
else:
nSketchClasses = FLAGS.nSketchClasses
print('Not Ignore BG;', nSketchClasses, 'classes')
data_aug = FLAGS.data_aug if mode == 'train' else False
model = adapted_deeplab_model.DeepLab(num_classes=nSketchClasses,
lrn_rate=FLAGS.learning_rate,
lrn_rate_end=FLAGS.learning_rate_end,
optimizer=FLAGS.optimizer,
upsample_mode=FLAGS.upsample_mode,
data_aug=data_aug,
image_down_scaling=FLAGS.image_down_scaling,
ignore_class_bg=FLAGS.ignore_class_bg,
mode=mode)
snapshot_saver = tf.train.Saver(max_to_keep=5)
tfconfig = tf.ConfigProto()
tfconfig.gpu_options.allow_growth = True
sess = tf.Session(config=tfconfig)
sess.run(tf.global_variables_initializer())
snapshot_dir = os.path.join(FLAGS.outputs_base_dir, FLAGS.snapshot_folder_name)
os.makedirs(snapshot_dir, exist_ok=True)
ckpt = tf.train.get_checkpoint_state(snapshot_dir)
if not ckpt:
if mode == 'train':
pretrained_model = FLAGS.resnet_pretrained_model_path
load_var = {var.op.name: var for var in tf.global_variables()
if var.op.name.startswith('ResNet')
and 'factor' not in var.op.name
and 'Adam' not in var.op.name
and 'beta1_power' not in var.op.name
and 'beta2_power' not in var.op.name
and 'fc_final_sketch46' not in var.op.name
and 'global_step' not in var.op.name # count from 0
}
snapshot_loader = tf.train.Saver(load_var)
print('Firstly training, loaded', pretrained_model)
snapshot_loader.restore(sess, pretrained_model)
else:
raise Exception("No pre-trained model for %s" % mode)
else:
load_var = {var.op.name: var for var in tf.global_variables()
if var.op.name.startswith('ResNet')
and 'global_step' not in var.op.name # count from 0
}
snapshot_loader = tf.train.Saver(load_var)
print('Trained model found, loaded', ckpt.model_checkpoint_path)
snapshot_loader.restore(sess, ckpt.model_checkpoint_path)
if mode == 'train':
log_dir = os.path.join(FLAGS.outputs_base_dir, FLAGS.log_folder_name)
os.makedirs(log_dir, exist_ok=True)
snapshot_file = os.path.join(snapshot_dir, 'iter_%d.tfmodel')
summary_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(log_dir, graph=sess.graph)
duration_time_n_step = 0
for n_iter in range(FLAGS.max_iteration):
start_time = time.time()
print('\n#' + str(n_iter))
## select image index
image_idx = random.randint(1, FLAGS.nTrainImgs)
## load images
image_name = 'L0_sample' + str(image_idx) + '.png' # e.g. L0_sample5564.png
# print("Load:", image_name)
image_path = os.path.join(FLAGS.data_base_dir, mode, 'DRAWING_GT', image_name)
im = load_image(image_path, mu) # shape = [1, H, W, 3]
# print("Ori shape", im.shape)
## load label
label_name = 'sample_' + str(image_idx) + '_class.mat' # e.g. sample_1_class.mat
label_path = os.path.join(FLAGS.data_base_dir, mode, 'CLASS_GT', label_name)
label = load_label(label_path) # shape = [1, H, W], [0, 46]
if FLAGS.ignore_class_bg:
label = label - 1 # [-1, 45]
label[label == -1] = 255 # [0-45, 255]
feed_dict = {model.images: im, model.labels: label}
_, learning_rate_, global_step, cost, pred, pred_label = \
sess.run([model.train_step,
model.learning_rate,
model.global_step,
model.cost,
model.pred,
model.pred_label],
feed_dict=feed_dict)
# print('pred.shape', pred.shape) # (1, H_scale, W_scale, nClasses)
print('learning_rate_', learning_rate_)
# print('global_step', global_step)
print('cost', cost)
## display left time
duration_time = time.time() - start_time
duration_time_n_step += duration_time
if n_iter % FLAGS.count_left_time_freq == 0 and n_iter != 0:
left_step = FLAGS.max_iteration - n_iter
left_sec = left_step / FLAGS.count_left_time_freq * duration_time_n_step
print("Duration_time_%d_step: %s. Left time: %s" % (
FLAGS.count_left_time_freq,
str(timedelta(seconds=duration_time_n_step)),
str(timedelta(seconds=left_sec))))
duration_time_n_step = 0
## summary
if n_iter % FLAGS.summary_write_freq == 0 and n_iter != 0:
summary_str = sess.run(summary_op, feed_dict=feed_dict)
summary_writer.add_summary(summary_str, n_iter)
summary_writer.flush()
## save model
if (n_iter + 1) % FLAGS.save_model_freq == 0 or (n_iter + 1) >= FLAGS.max_iteration:
snapshot_saver.save(sess, snapshot_file % (n_iter + 1))
print('model saved to ' + snapshot_file % (n_iter + 1))
print('Training done.')
elif mode == 'val' or mode == 'test':
def fast_hist(a, b, n):
"""
:param a: gt
:param b: pred
"""
k = (a >= 0) & (a < n)
return np.bincount(n * a[k].astype(int) + b[k], minlength=n ** 2).reshape(n, n)
use_dcrf = kwargs['use_dcrf']
eval_base_dir = os.path.join(FLAGS.outputs_base_dir, 'eval_results')
os.makedirs(eval_base_dir, exist_ok=True)
nImgs = FLAGS.nTestImgs if mode == 'test' else FLAGS.nValImgs
colorMap = scipy.io.loadmat(os.path.join(FLAGS.data_base_dir, 'colorMapC46.mat'))['colorMap']
outstr = mode + ' mode\n'
cat_max_len = 16
hist = np.zeros((FLAGS.nSketchClasses, FLAGS.nSketchClasses))
for imgIndex in range(1, nImgs + 1):
## load images
image_name = 'L0_sample' + str(imgIndex) + '.png' # e.g. L0_sample5564.png
image_path = os.path.join(FLAGS.data_base_dir, mode, 'DRAWING_GT', image_name)
test_image = load_image(image_path, mu) # shape = [1, H, W, 3]
## load gt_label
label_name = 'sample_' + str(imgIndex) + '_class.mat' # e.g. sample_1_class.mat
label_path = os.path.join(FLAGS.data_base_dir, mode, 'CLASS_GT', label_name)
gt_label = load_label(label_path) # shape = [1, H, W]
print('#' + str(imgIndex) + '/' + str(nImgs) + ': ' + image_path)
feed_dict = {model.images: test_image, model.labels: 0}
pred, pred_label_no_crf = sess.run([model.pred, model.pred_label], feed_dict=feed_dict)
if FLAGS.ignore_class_bg:
pred_label_no_crf = pred_label_no_crf + 1 # [1, 46]
# print('@ pred.shape ', pred.shape) # (1, H, W, nSketchClasses)
# print(pred_label_no_crf.shape) # shape = [1, H, W, 1]
if use_dcrf:
prob_arr = np.squeeze(pred)
prob_arr = prob_arr.transpose((2, 0, 1)) # shape = (nSketchClasses, H, W)
d_image = np.array(np.squeeze(test_image), dtype=np.uint8) # shape = (H, W, 3)
pred_label_crf = seg_densecrf(prob_arr, d_image, nSketchClasses) # shape=[H, W]
if FLAGS.ignore_class_bg:
pred_label_crf = pred_label_crf + 1 # [1, 46]
hist += fast_hist(np.squeeze(gt_label).flatten(),
pred_label_crf.flatten(),
FLAGS.nSketchClasses)
else:
hist += fast_hist(np.squeeze(gt_label).flatten(),
np.squeeze(pred_label_no_crf).flatten(),
FLAGS.nSketchClasses)
if imgIndex == nImgs:
## ignore bg pixel with value 0
if FLAGS.ignore_class_bg:
hist = hist[1:, 1:]
if use_dcrf:
print('\nRound', str(imgIndex), ', Use CRF')
outstr += '\nRound: ' + str(imgIndex) + ', Use CRF' + '\n'
else:
print('\nRound', str(imgIndex), ', Not Use CRF')
outstr += '\nRound: ' + str(imgIndex) + ', Not Use CRF' + '\n'
# overall accuracy
acc = np.diag(hist).sum() / hist.sum()
print('>>> overall accuracy', acc)
outstr += '>>> overall accuracy ' + str(acc) + '\n'
# mAcc
acc = np.diag(hist) / hist.sum(1)
mean_acc = np.nanmean(acc)
print('>>> mean accuracy', mean_acc)
outstr += '>>> mean accuracy ' + str(mean_acc) + '\n'
# mIoU
iou = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
mean_iou = np.nanmean(iou)
print('>>> mean IoU', mean_iou)
outstr += '>>> mean IoU ' + str(mean_iou) + '\n'
# FWIoU
freq = hist.sum(1) / hist.sum()
fw_iou = (freq[freq > 0] * iou[freq > 0]).sum()
print('>>> freq weighted IoU', fw_iou)
print('\n')
outstr += '>>> freq weighted IoU ' + str(fw_iou) + '\n'
# IoU of each class
print('>>> IoU of each class')
outstr += '\n>>> IoU of each class' + '\n'
for classIdx in range(nSketchClasses):
if FLAGS.ignore_class_bg:
cat_name = colorMap[classIdx][0][0]
else:
if classIdx == 0:
cat_name = 'background'
else:
cat_name = colorMap[classIdx - 1][0][0]
singlestr = ' >>> '
cat_len = len(cat_name)
pad = ''
for ipad in range(cat_max_len - cat_len):
pad += ' '
singlestr += cat_name + pad + str(iou[classIdx])
print(singlestr)
outstr += singlestr + '\n'
# write validation result to txt
write_path = os.path.join(eval_base_dir, mode + '_results.txt')
fp = open(write_path, 'a')
fp.write(outstr)
fp.close()
else: # 'inference'
inference_ids = [kwargs['inference_id']]
inference_dataset = kwargs['inference_dataset']
black_bg = kwargs['black_bg']
use_dcrf = kwargs['use_dcrf']
colorMap = scipy.io.loadmat(os.path.join(FLAGS.data_base_dir, 'colorMapC46.mat'))['colorMap']
infer_result_base_dir = os.path.join(FLAGS.outputs_base_dir, 'inference_results', inference_dataset)
os.makedirs(infer_result_base_dir, exist_ok=True)
for img_count, img_id in enumerate(inference_ids):
image_name = 'L0_sample' + str(img_id) + '.png' # e.g. L0_sample5564.png
image_path = os.path.join(FLAGS.data_base_dir, inference_dataset, 'DRAWING_GT', image_name)
infer_image, infer_image_raw = load_image(image_path, mu, return_raw=True) # shape = [1, H, W, 3] / [H, W, 3]
print('\n#' + str(img_count + 1) + '/' + str(len(inference_ids)) + ': ' + image_name)
feed_dict = {model.images: infer_image, model.labels: 0}
pred, pred_label_no_crf, feat_visual \
= sess.run([model.pred, model.pred_label, model.feat_visual], feed_dict=feed_dict)
print('@ pred.shape ', pred.shape) # (1, H, W, nSketchClasses)
print('@ pred_label_no_crf.shape ', pred_label_no_crf.shape) # shape = [1, H, W, 1], contains [0, nClasses)
# print('@ feat_visual.shape ', feat_visual.shape) # shape = (1, 94, 94, 512)
if use_dcrf:
prob_arr = np.squeeze(pred)
prob_arr = prob_arr.transpose((2, 0, 1)) # shape = (nSketchClasses, H, W)
d_image = np.array(np.squeeze(infer_image), dtype=np.uint8) # shape = (H, W, 3)
pred_label_crf = seg_densecrf(prob_arr, d_image, nSketchClasses) # shape=[H, W], contains [0-46/47]
save_base_dir_crf = os.path.join(infer_result_base_dir, 'deeplab_output_crf')
os.makedirs(save_base_dir_crf, exist_ok=True)
if FLAGS.ignore_class_bg:
pred_label_crf += 1
pred_label_crf[infer_image_raw[:, :, 0] != 0] = 0 # [H, W]
save_path_crf = os.path.join(save_base_dir_crf, 'sem_result_' + str(img_id) + '.png')
visualize_semantic_segmentation(pred_label_crf, colorMap, black_bg=black_bg, save_path=save_path_crf)
else:
save_base_dir_no_crf = os.path.join(infer_result_base_dir, 'deeplab_output_no_crf')
os.makedirs(save_base_dir_no_crf, exist_ok=True)
if FLAGS.ignore_class_bg:
pred_label_no_crf += 1
pred_label_no_crf = np.squeeze(pred_label_no_crf)
pred_label_no_crf[infer_image_raw[:, :, 0] != 0] = 0 # [H, W]
save_path_no_crf = os.path.join(save_base_dir_no_crf, 'sem_result_' + str(img_id) + '.png')
visualize_semantic_segmentation(pred_label_no_crf, colorMap, black_bg=black_bg, save_path=save_path_no_crf)
def train_model(data_folder,
data_name,
level,
model_name,
is_aspect_term=True,
classWeights=None,
imBalanced=False):
config.data_folder = data_folder
config.data_name = data_name
if not os.path.exists(os.path.join(config.checkpoint_dir, data_folder)):
os.makedirs(os.path.join(config.checkpoint_dir, data_folder))
config.level = level
config.model_name = model_name
config.is_aspect_term = is_aspect_term
config.init_input()
config.exp_name = '{}_{}_wv_{}'.format(model_name, level,
config.word_embed_type)
config.exp_name = config.exp_name + '_update' if config.word_embed_trainable else config.exp_name + '_fix'
if config.use_aspect_input:
config.exp_name += '_aspv_{}'.format(config.aspect_embed_type)
config.exp_name = config.exp_name + '_update' if config.aspect_embed_trainable else config.exp_name + '_fix'
if config.use_elmo:
config.exp_name += '_elmo_alone_{}_mode_{}_{}'.format(
config.use_elmo_alone, config.elmo_output_mode,
'update' if config.elmo_trainable else 'fix')
print(config.exp_name)
model = SentimentModel(config)
test_input = load_input_data(
data_folder, 'test', level, config.use_text_input,
config.use_text_input_l, config.use_text_input_r,
config.use_text_input_r_with_pad, config.use_aspect_input,
config.use_aspect_text_input, config.use_loc_input,
config.use_offset_input, config.use_mask)
test_label = load_label(data_folder, 'test')
if not os.path.exists(
os.path.join(config.checkpoint_dir, '%s/%s.hdf5' %
(data_folder, config.exp_name))):
start_time = time.time()
train_input = load_input_data(
data_folder, 'train', level, config.use_text_input,
config.use_text_input_l, config.use_text_input_r,
config.use_text_input_r_with_pad, config.use_aspect_input,
config.use_aspect_text_input, config.use_loc_input,
config.use_offset_input, config.use_mask)
train_label = load_label(data_folder, 'train')
valid_input = load_input_data(
data_folder, 'valid', level, config.use_text_input,
config.use_text_input_l, config.use_text_input_r,
config.use_text_input_r_with_pad, config.use_aspect_input,
config.use_aspect_text_input, config.use_loc_input,
config.use_offset_input, config.use_mask)
valid_label = load_label(data_folder, 'valid')
'''
Note: Here I combine the training data and validation data together, use them as training input to the model,
while I use test data to server as validation input. The reason behind is that i want to fully explore how
well can the model perform on the test data (Keras's ModelCheckpoint callback can help usesave the model
which perform best on validation data (here the test data)).
But generally, we won't do that, because test data will not (and should not) be accessible during training
process.
'''
# train_combine_valid_input = []
# for i in range(len(train_input)):
# train_combine_valid_input.append(train_input[i] + valid_input[i])
# train_combine_valid_label = train_label + valid_label
model.train(train_input, train_label, test_input, test_label,
classWeights, imBalanced)
# model.train(train_combine_valid_input, train_combine_valid_label, test_input, test_label)
elapsed_time = time.time() - start_time
print('training time:',
time.strftime("%H:%M:%S", time.gmtime(elapsed_time)))
# load the best model
model.load()
# print('score over valid data...')
# model.score(valid_input, valid_label)
print('score over test data...')
model.score(test_input, test_label)
import torch.nn as nn
import data_loader
import iter_utils
import torch.utils.data as data_utils
from models import *
## load data
data_arr_01 = data_loader.load_data('data/pgb/SF01/vib_data_1.txt')
# data_arr_03 = data_loader.load_data('data/pgb/SF03/vib_data_1.txt')
# data_arr_01 = data_loader.resample_arr(data_arr_01, num=240) # add for Ince's model
# data_arr_03 = data_loader.resample_arr(data_arr_03, num=240) # add for Ince's model
# data_arr_01, _ = data_loader.fft_arr(data_arr_01) # add for fft wdcnn
# data_arr_03, _ = data_loader.fft_arr(data_arr_03) # add for fft wdcnn
# data_arr_01 = data_loader.stft_arr(data_arr_01) # add for stft-LeNet
# data_arr_03 = data_loader.stft_arr(data_arr_03)
label_vec = data_loader.load_label('data/pgb/SF01/label_vec.txt')
trainset_01, testset_01 = data_loader.split_set(data_arr_01, label_vec)
# trainset_03, testset_03 = data_loader.split_set(data_arr_03, label_vec)
train_loader = data_utils.DataLoader(dataset=trainset_01,
batch_size=512,
shuffle=True,
num_workers=2)
test_loader = data_utils.DataLoader(dataset=testset_01,
batch_size=512,
shuffle=True,
num_workers=2)
print('Number of training samples: {}'.format(len(train_loader.dataset)))
print('Number of testing samples: {}'.format(len(test_loader.dataset)))
print()
def __init__(self):
self.estimator = self._make_estimator()
self.vocab = data_loader.load_vocab()
self.tag_dict = data_loader.load_tag()
self.label_dict = data_loader.load_label()
def semantic_main(**kwargs):
mode = str(kwargs['mode'])
ignore_class_bg = kwargs['ignore_class_bg']
config = Config()
print('Mode:', mode, '; ignore_class_bg:', ignore_class_bg)
snapshot_dir = os.path.join(config.outputs_base_dir,
config.snapshot_folder_name)
os.makedirs(snapshot_dir, exist_ok=True)
model_fname = os.path.join(snapshot_dir,
config.model_fname_start + '%d.pth')
# BG ignoring
if ignore_class_bg:
nSketchClasses = config.nSketchClasses - 1
else:
nSketchClasses = config.nSketchClasses
colorMap = scipy.io.loadmat(
os.path.join(config.data_base_dir, 'colorMapC46.mat'))['colorMap']
use_gpu = torch.cuda.is_available()
data_transforms = transforms.Compose([
transforms.ToTensor(), # [3, H, W], [0.0-1.0]
])
model = getattr(adapted_deeplab_pytorch,
'resnet101')(num_classes=nSketchClasses,
up_mode=config.upsample_mode)
if mode == 'train':
init_with = kwargs['init_with']
log_info = kwargs['log_info']
model.train()
start_iter = 0
# load weight
if init_with == 'resnet':
# load weight only in resnet-trained model
valid_state_dict = model.state_dict()
trained_state_dict = torch.load(
config.resnet_pretrained_model_path)
for var_name in model.state_dict():
if var_name in trained_state_dict:
valid_state_dict[var_name] = trained_state_dict[var_name]
model.load_state_dict(valid_state_dict)
print('Loaded', config.resnet_pretrained_model_path)
elif init_with == 'last':
# find the last trained model
last_model_path, max_iter = find_last_model(
snapshot_dir, config.model_fname_start)
model.load_state_dict(torch.load(last_model_path))
print('Loaded', last_model_path)
start_iter = max_iter
else: # init_with == 'none'
print('Training from fresh start.')
if log_info:
from logger import Logger
log_dir = os.path.join(config.outputs_base_dir,
config.log_folder_name)
os.makedirs(log_dir, exist_ok=True)
logger = Logger(log_dir)
if use_gpu:
model = model.cuda()
criterion = nn.CrossEntropyLoss().cuda()
opt_param_groups = [{
'params': model.conv1.parameters()
}, {
'params': model.bn1.parameters()
}, {
'params': model.layer1.parameters()
}, {
'params': model.layer2.parameters()
}, {
'params': model.layer3.parameters()
}, {
'params': model.layer4.parameters()
}, {
'params':
iter([
model.fc1_sketch46_c0.weight, model.fc1_sketch46_c1.weight,
model.fc1_sketch46_c2.weight, model.fc1_sketch46_c3.weight
])
}, {
'params':
iter([
model.fc1_sketch46_c0.bias, model.fc1_sketch46_c1.bias,
model.fc1_sketch46_c2.bias, model.fc1_sketch46_c3.bias
]),
'weight_decay':
0.
}]
if config.upsample_mode == 'deconv':
opt_param_groups.append({'params': model.deconv.parameters()})
if config.optimizer == 'sgd':
optimizer = optim.SGD(opt_param_groups,
lr=config.base_lr,
momentum=0.9,
weight_decay=0.0005)
elif config.optimizer == 'adam':
optimizer = optim.Adam(opt_param_groups,
lr=config.base_lr,
weight_decay=0.0005)
else:
raise NameError("Unknown optimizer type %s!" % config.optimizer)
losses = AverageMeter()
duration_time_n_step = 0
for n_iter in range(start_iter, config.max_iteration):
start_time = time.time()
# set lr to the same to all layers
lr = (config.base_lr - config.end_lr) * \
math.pow(1 - float(min(n_iter, config.end_decay_step)) / config.end_decay_step, 0.9) + config.end_lr
if config.multiplied_lr:
for g in range(6):
optimizer.param_groups[g]['lr'] = lr
optimizer.param_groups[6]['lr'] = lr * 10
optimizer.param_groups[7]['lr'] = lr * 20
if config.upsample_mode == 'deconv':
optimizer.param_groups[8]['lr'] = lr
else:
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# select image index
image_idx = random.randint(1, config.nTrainImgs)
# load images
image_name = 'L0_sample' + str(
image_idx) + '.png' # e.g. L0_sample5564.png
# print("Load:", image_name)
image_path = os.path.join(config.data_base_dir, mode, 'DRAWING_GT',
image_name)
im = datasets.folder.default_loader(
image_path) # [H, W, 3], [0-255]
inputs = data_transforms(im) # [3, H, W], ([0.0-1.0]-mean)/std
if use_gpu:
inputs = Variable(inputs.cuda())
else:
inputs = Variable(inputs)
# sketchyscene change: we need to scale the inputs to ([0.0-255.0]-mean),
# following the TF model
inputs = inputs.permute(1, 2, 0) # [H, W, 3]
inputs = inputs * 255.0
for i in range(3):
inputs[:, :, i] = inputs[:, :, i] - config.mean[i]
inputs = inputs.permute(2, 0, 1) # [3, H, W]
# load gt_label
label_name = 'sample_' + str(
image_idx) + '_class.mat' # e.g. sample_1_class.mat
label_path = os.path.join(config.data_base_dir, mode, 'CLASS_GT',
label_name)
gt_label = load_label(label_path) # shape = [1, H, W], [0, 46]
if ignore_class_bg:
gt_label = gt_label - 1 # [-1, 45]
gt_label[gt_label == -1] = 255 # [0-45, 255]
h, w = gt_label.shape[1], gt_label.shape[2]
gt_label = torch.LongTensor(gt_label) # [1, H, W]
if use_gpu:
gt_label = Variable(gt_label.cuda())
else:
gt_label = Variable(gt_label)
outputs = model(inputs.unsqueeze(0))
if config.upsample_mode == 'deconv':
outputs_up = outputs
else:
outputs_up = nn.UpsamplingBilinear2d(
(h, w))(outputs) # [1, nClasses, H, W]
# calculate loss
# ignore illegal labels
gt_label_flatten = gt_label.view(-1, ) # [H * W]
outputs_up_flatten = torch.t(
outputs_up.view(outputs_up.shape[1], -1)) # [H * W, nClasses]
if ignore_class_bg: # [0-45, 255], ignore 255
mask = torch.lt(gt_label_flatten,
nSketchClasses) # lower than 46, [H * W]
gt_label_flatten = gt_label_flatten[mask] # [<= H * W]
outputs_up_flatten = outputs_up_flatten[
mask, :] # [<= H * W, nClasses]
# outputs_up_flatten: [N, C], gt_label_flatten: [N]
loss = criterion(outputs_up_flatten, gt_label_flatten)
losses.update(loss.data[0], 1)
# optimization
optimizer.zero_grad()
loss.backward()
optimizer.step()
# display loss
print("Step: {0}. Lr: {1}. Loss: {loss.val:.6f} ({loss.avg:.4f})".
format(n_iter + 1, lr, loss=losses))
# display left time
duration_time = time.time() - start_time
duration_time_n_step += duration_time
if n_iter % config.count_left_time_freq == 0 and n_iter != 0:
left_step = config.max_iteration - n_iter
left_sec = left_step / config.count_left_time_freq * duration_time_n_step
print("Left time: {}".format(str(timedelta(seconds=left_sec))))
duration_time_n_step = 0
# summary
if log_info and n_iter % config.summary_write_freq == 0 and n_iter != 0:
info = {'loss': loss.item(), 'learning_rate': lr}
for tag, value in info.items():
# print('tag, value', tag, value)
logger.scalar_summary(tag, value, n_iter + 1)
# save model
if (n_iter + 1) % config.save_model_freq == 0 or (
n_iter + 1) >= config.max_iteration:
torch.save(model.state_dict(), model_fname % (n_iter + 1))
print('model saved to ' + model_fname % (n_iter + 1))
print('Training done.')
elif mode == 'val' or mode == 'test':
dcrf = kwargs['dcrf']
def fast_hist(a, b, n):
"""
:param a: gt
:param b: pred
"""
k = (a >= 0) & (a < n)
return np.bincount(n * a[k].astype(int) + b[k],
minlength=n**2).reshape(n, n)
eval_base_dir = os.path.join(config.outputs_base_dir,
config.eval_folder_name)
os.makedirs(eval_base_dir, exist_ok=True)
model_path, _ = find_last_model(snapshot_dir, config.model_fname_start)
model.eval()
model.load_state_dict(torch.load(model_path))
print('Load', model_path)
if use_gpu:
model = model.cuda()
nImgs = config.nTestImgs if mode == 'test' else config.nValImgs
outstr = mode + ' mode\n'
cat_max_len = 16
hist = np.zeros((config.nSketchClasses, config.nSketchClasses))
for imgIndex in range(1, nImgs + 1):
# load images
image_name = 'L0_sample' + str(
imgIndex) + '.png' # e.g. L0_sample5564.png
image_path = os.path.join(config.data_base_dir, mode, 'DRAWING_GT',
image_name)
im = datasets.folder.default_loader(
image_path) # [H, W, 3], [0-255]
im_np = np.array(im, dtype=np.uint8)
h, w = np.shape(im)[0], np.shape(im)[1]
inputs = data_transforms(im) # [3, H, W], [0.0-1.0]
if use_gpu:
inputs = Variable(inputs.cuda())
else:
inputs = Variable(inputs)
# sketchyscene change: we need to scale the inputs to ([0.0-255.0]-mean),
# following the TF model
inputs = inputs.permute(1, 2, 0) # [H, W, 3]
inputs = inputs * 255.0
for i in range(3):
inputs[:, :, i] = inputs[:, :, i] - config.mean[i]
inputs = inputs.permute(2, 0, 1) # [3, H, W]
outputs = model(inputs.unsqueeze(0))
if config.upsample_mode == 'deconv':
outputs_up = outputs
else:
outputs_up = nn.UpsamplingBilinear2d(
(h, w))(outputs) # [1, nClasses, H, W]
# load gt_label
label_name = 'sample_' + str(
imgIndex) + '_class.mat' # e.g. sample_1_class.mat
label_path = os.path.join(config.data_base_dir, mode, 'CLASS_GT',
label_name)
gt_label = load_label(label_path) # shape = [1, H, W], [0, 46]
print('#' + str(imgIndex) + '/' + str(nImgs) + ': ' + image_path)
if dcrf:
prob_arr = outputs_up.data.cpu().numpy().squeeze(
0) # [nClasses, H, W]
pred_label = seg_densecrf(
prob_arr, im_np,
nSketchClasses) # shape=[H, W], contains [0-45/46]
else:
_, pred = torch.max(outputs_up, 1)
pred_label = pred.data.cpu().numpy() # [1, H, W]
pred_label = pred_label.squeeze().astype(
np.uint8) # [H, W], [0-45/46]
if ignore_class_bg:
pred_label = pred_label + 1 # [1-46]
hist += fast_hist(
np.squeeze(gt_label).flatten(), pred_label.flatten(),
config.nSketchClasses)
# ignore bg pixel with value 0
if ignore_class_bg:
hist = hist[1:, 1:]
if dcrf:
print('\nRound', str(nImgs), ', Use CRF')
outstr += '\nRound: ' + str(nImgs) + ', Use CRF' + '\n'
else:
print('\nRound', str(nImgs), ', Not Use CRF')
outstr += '\nRound: ' + str(nImgs) + ', Not Use CRF' + '\n'
# overall accuracy
acc = np.diag(hist).sum() / hist.sum()
print('>>> overall accuracy', acc)
outstr += '>>> overall accuracy ' + str(acc) + '\n'
# mAcc
acc = np.diag(hist) / hist.sum(1)
mean_acc = np.nanmean(acc)
print('>>> mean accuracy', mean_acc)
outstr += '>>> mean accuracy ' + str(mean_acc) + '\n'
# mIoU
iou = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
mean_iou = np.nanmean(iou)
print('>>> mean IoU', mean_iou)
outstr += '>>> mean IoU ' + str(mean_iou) + '\n'
# FWIoU
freq = hist.sum(1) / hist.sum()
fw_iou = (freq[freq > 0] * iou[freq > 0]).sum()
print('>>> freq weighted IoU', fw_iou)
print('\n')
outstr += '>>> freq weighted IoU ' + str(fw_iou) + '\n'
# IoU of each class
print('>>> IoU of each class')
outstr += '\n>>> IoU of each class' + '\n'
for classIdx in range(nSketchClasses):
if ignore_class_bg:
cat_name = colorMap[classIdx][0][0]
else:
if classIdx == 0:
cat_name = 'background'
else:
cat_name = colorMap[classIdx - 1][0][0]
singlestr = ' >>> '
cat_len = len(cat_name)
pad = ''
for ipad in range(cat_max_len - cat_len):
pad += ' '
singlestr += cat_name + pad + str(iou[classIdx])
print(singlestr)
outstr += singlestr + '\n'
# write validation result to txt
write_path = os.path.join(eval_base_dir, mode + '_results.txt')
fp = open(write_path, 'a')
fp.write(outstr)
fp.close()
elif mode == 'inference':
dcrf = kwargs['dcrf']
image_idx = kwargs['inference_id']
dataset_type = kwargs['inference_dataset']
black_bg = kwargs['black_bg']
infer_result_base_dir = os.path.join(config.outputs_base_dir,
config.inference_folder_name,
dataset_type)
os.makedirs(infer_result_base_dir, exist_ok=True)
model_path, _ = find_last_model(snapshot_dir, config.model_fname_start)
model.eval()
model.load_state_dict(torch.load(model_path))
print('Load', model_path)
if use_gpu:
model = model.cuda()
# load images
image_name = 'L0_sample' + str(
image_idx) + '.png' # e.g. L0_sample5564.png
print('image_name', image_name)
image_path = os.path.join(config.data_base_dir, dataset_type,
'DRAWING_GT', image_name)
im = datasets.folder.default_loader(image_path) # [H, W, 3], [0-255]
im_np = np.array(im, dtype=np.uint8)
h, w = np.shape(im)[0], np.shape(im)[1]
inputs = data_transforms(im) # [3, H, W], [0.0-1.0]
if use_gpu:
inputs = Variable(inputs.cuda())
else:
inputs = Variable(inputs)
# sketchyscene change: we need to scale the inputs to ([0.0-255.0]-mean),
# following the TF model
inputs = inputs.permute(1, 2, 0) # [H, W, 3]
inputs = inputs * 255.0
for i in range(3):
inputs[:, :, i] = inputs[:, :, i] - config.mean[i]
inputs = inputs.permute(2, 0, 1) # [3, H, W]
outputs = model(inputs.unsqueeze(0))
if config.upsample_mode == 'deconv':
outputs_up = outputs
else:
outputs_up = nn.UpsamplingBilinear2d(
(h, w))(outputs) # [1, nClasses, H, W]
if dcrf:
prob_arr = outputs_up.data.cpu().numpy().squeeze(
0) # [nClasses, H, W]
pred_label = seg_densecrf(
prob_arr, im_np,
nSketchClasses) # shape=[H, W], contains [0-45/46]
else:
_, pred = torch.max(outputs_up, 1)
pred_label = pred.data.cpu().numpy() # [1, H, W]
pred_label = pred_label.squeeze().astype(
np.uint8) # [H, W], [0-45/46]
if ignore_class_bg:
pred_label = pred_label + 1 # [1-46]
# filter with binary mask
pred_label[im_np[:, :, 0] != 0] = 0 # [H, W], [0-46]
subdir = 'deeplab_output_crf' if dcrf else 'deeplab_output_no_crf'
save_base_dir = os.path.join(infer_result_base_dir, subdir)
os.makedirs(save_base_dir, exist_ok=True)
save_path = os.path.join(save_base_dir,
'sem_result_' + str(image_idx) + '.png')
visualize_semantic_segmentation(pred_label,
colorMap,
black_bg=black_bg,
save_path=save_path)
def train_model(data_folder,
data_name,
level,
model_name,
is_aspect_term=True):
config.data_folder = data_folder
config.data_name = data_name
# 新建存处
if not os.path.exists(os.path.join(config.checkpoint_dir, data_folder)):
os.makedirs(os.path.join(config.checkpoint_dir, data_folder))
config.level = level # char 中文
config.model_name = model_name # atae_lstm or tsa
config.is_aspect_term = is_aspect_term # true
config.init_input()
# 给保存时候的名字
config.exp_name = '{}_{}_wv_{}'.format(model_name, level,
config.word_embed_type)
# 可更新
config.exp_name = config.exp_name + '_update' if config.word_embed_trainable else config.exp_name + '_fix'
if config.use_aspect_input:
config.exp_name += '_aspv_{}'.format(config.aspect_embed_type)
config.exp_name = config.exp_name + '_update' if config.aspect_embed_trainable else config.exp_name + '_fix'
# 不用 ，否则tensorflow_hub问题难解决？
# if config.use_elmo:
# config.exp_name += '_elmo_alone_{}_mode_{}_{}'.format(config.use_elmo_alone, config.elmo_output_mode,
# 'update' if config.elmo_trainable else 'fix')
print(config.exp_name)
# 建
model = SentimentModel(config)
test_input = load_input_data(data_folder, 'test', level,
config.use_text_input,
config.use_aspect_input,
config.use_aspect_text_input)
test_label = load_label(data_folder, 'test')
print(test_input)
# there's no dev data of laptop
# dev_input = load_input_data(data_folder, 'valid', level, config.use_text_input,
# config.use_aspect_input,config.use_aspect_text_input)
#
# dev_label = load_label(data_folder, 'valid')
#
# print(dev_input)
# 无现有模型，开始训练
if not os.path.exists(
os.path.join(config.checkpoint_dir, '%s/%s.hdf5' %
(data_folder, config.exp_name))):
start_time = time.time()
train_input = load_input_data(data_folder, 'train', level,
config.use_text_input,
config.use_aspect_input,
config.use_aspect_text_input)
train_label = load_label(data_folder, 'train')
# valid_input = load_input_data(data_folder, 'valid', level, config.use_text_input,
# config.use_aspect_input, config.use_aspect_text_input)
# valid_label = load_label(data_folder, 'valid')
# train
model.train(train_input, train_label, test_input, test_label)
elapsed_time = time.time() - start_time
print('training time:',
time.strftime("%H:%M:%S", time.gmtime(elapsed_time)))
# load the best model
model.load()
print("start to score ...")
# print('score over dev data...')
# model.score(dev_input, dev_label)
print('score over test data...')
model.score(test_input, test_label)
print("score done!")
print('start to predict and save the results...')
# print('predict over dev data...')
# result = model.predict(dev_input)
# print("save prediction and actual labels of dev...")
# print(dev_label)
# print(result)
# concat(result, dev_label, model_name, 1, config.word_embed_type)
print('predict over test data...')
result2 = model.predict(test_input)
print("save prediction and actual labels of dev...")
concat(result2, test_label, model_name, 2, config.word_embed_type)
print('predict and save the results done!')
print('totally done!')