Esempio n. 1
0
def eval_regression(task):
    model_path = './model/{}.model'.format(task)

    D_m_text, D_m_audio, D_m_video, D_m_context = 300, 384, 35, 300
    D_g, D_p, D_e, D_h, D_a = 150, 150, 100, 100, 100

    cuda = torch.cuda.is_available()

    print('Loading model...')
    model = RegressionModel(D_m_text,
                            D_m_audio,
                            D_m_video,
                            D_m_context,
                            D_g,
                            D_p,
                            D_e,
                            D_h,
                            dropout_rec=0.1,
                            dropout=0.25)
    if cuda:
        model.cuda()
    model.load_state_dict(torch.load(model_path))

    loss_function = MaskedMSELoss()

    print('Evaluating model...')
    _, _, test_loader = train_regression.get_MOSEI_loaders(
        './data/regression.pkl', valid=0.0, batch_size=128, num_workers=0)

    _, mae, _, labels, preds, masks, sample_ids = train_regression.train_or_eval_model(
        model, loss_function, test_loader, None, cuda)

    # gather labels and predictions
    df = pd.DataFrame([(sample_id, label, pred)
                       for label, pred, mask, sample_id in zip(
                           labels, preds, masks, sample_ids) if mask == 1],
                      columns=['sample_id', 'label', 'pred'])
    df['diff'] = (df.label - df.pred).abs()
    df['label_class'] = df.label.apply(discretize)
    df['pred_class'] = df.pred.apply(discretize)
    df = df.sort_values(by='diff', ascending=False)

    if_correct = df.label_class == df.pred_class

    print('mae =', mean_absolute_error(df.label, df.pred))
    print('acc =', if_correct.sum() / len(if_correct))
    df.to_csv('./analysis/{}.csv'.format(task), index=False)
    else:
        print('Running on CPU')
    print("Tensorboard logs in " + args.log_dir)

    batch_size = args.batch_size
    n_classes  = 6
    cuda       = args.cuda
    n_epochs   = args.epochs
    D_m_text, D_m_audio, D_m_video, D_m_context = 300, 384, 35, 300
    D_g, D_p, D_e, D_h, D_a = 150, 150, 100, 100, 100

    # Instantiate model
    model = RegressionModel(D_m_text, D_m_audio, D_m_video, D_m_context, D_g, D_p, D_e, D_h, dropout_rec=args.rec_dropout, dropout=args.dropout)

    if cuda:
        model.cuda()
    loss_function = MaskedMSELoss()

    # Get optimizer and relevant dataloaders
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2)
    train_loader, valid_loader, test_loader = get_MOSEI_loaders('./data/regression.pkl', valid=0.0, batch_size=batch_size, num_workers=0)
    best_loss, best_label, best_pred, best_mask, best_pear = None, None, None, None, None

    # Training loop
    for e in tqdm(range(n_epochs), desc = 'MOSEI Regression'):
        train_loss, train_mae, train_pear,_,_,_ = train_or_eval_model(model, loss_function, train_loader, e, optimizer, True)
        test_loss, test_mae, test_pear, test_label, test_pred, test_mask = train_or_eval_model(model, loss_function, test_loader, e)
        writer.add_scalar("Train Loss - MOSEI Regression", train_loss, e)
        writer.add_scalar("Test Loss - MOSEI Regression", test_loss, e)
        writer.add_scalar("Train MAE - MOSEI Regression", train_mae, e)
        writer.add_scalar("Test MAE - MOSEI Regression", test_mae, e)