Пример #1
0
    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()
Пример #2
0
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))
Пример #3
0
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)
Пример #4
0
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)
Пример #5
0
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)
Пример #6
0
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()
Пример #7
0
 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()
Пример #8
0
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)
Пример #9
0
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!')