# Power Analysis

```python
import cPickle
import datetime
import glob
import gzip
import os
import random
import re
import shutil
import subprocess
import time
import uuid

import cdpybio as cpb
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
pd.options.mode.chained_assignment = None  # default='warn'
import pybedtools as pbt
import scipy
import scipy.stats as stats
import seaborn as sns
import statsmodels.api as sm
import vcf as pyvcf

import cardipspy as cpy
import ciepy

%matplotlib inline
%load_ext rpy2.ipython

dy_name = 'power_analysis'

import socket
if socket.gethostname() == 'fl-hn1' or socket.gethostname() == 'fl-hn2':
    dy = os.path.join(ciepy.root, 'sandbox', 'tmp', dy_name)
    cpy.makedir(dy)
    pbt.set_tempdir(dy)

outdir = os.path.join(ciepy.root, 'output', dy_name)
cpy.makedir(outdir)

private_outdir = os.path.join(ciepy.root, 'private_output', dy_name)
cpy.makedir(private_outdir)
```

    The rpy2.ipython extension is already loaded. To reload it, use:
      %reload_ext rpy2.ipython

```python
exp = pd.read_table(os.path.join(ciepy.root, 'output', 'eqtl_input',
                                 'tpm_log_filtered_phe_std_norm_peer_resid.tsv'), index_col=0)

rna_meta = pd.read_table(os.path.join(ciepy.root, 'output', 'input_data', 'rnaseq_metadata.tsv'),
                         index_col=0)

fn = os.path.join(ciepy.root, 'output/eqtl_processing/eqtls01/qvalues.tsv')
qvalues = pd.read_table(fn, index_col=0)
```

```python
def power_plots(tdf):
    for i in np.arange(40, 110, 10):
        df = pd.crosstab(tdf['power_{}_bin'.format(i)], tdf['sig_{}'.format(i)])
        df['percent'] = df[True] / df.sum(axis=1)
        plt.figure()
        df.percent.plot()
        plt.title(i)
        plt.ylabel('Observed power')
        plt.xlabel('Estimated power');
```

I want to look at power a few different ways.
I want to take my unrelateds and subsample them like GTEx does to see how many
eGenes I gain as I increase my sample size. I should probably run the full eQTL
analysis (with permutations) for the full unrelated set as well.

I'm also going to make an estimate of
my power to detect GTEx and Geuvadis eQTLs and compare them to my actual ability
to detect the eQTLs. For each Geuvadis/GTEx eQTL, I'll calculate an $R^2$. Since
Geuvadis provides $r$ for each eQTL, I can just square $r$. GTEx provides $\beta$,
so I can calculate $R^2$ for a given association as

\begin{align}
R^2 = \left(\beta \frac{s_v}{s_g}\right)^2,
\end{align}

where $s_v$ is the standard deviation of my genotypes at the position and $s_g$ is the
standard deviation of the gene's expression in the GTEx data. I can then calculate

\begin{align}
f^2 = \frac{R^2}{1 - R^2},
\end{align}

which I can use in a power test for linear regression.

The power calculation is kind of slow when done so many times so this notebook takes a while to run
if you need to redo the power calculations.

### Derivation of $R^2$ estimate

For simple linear regression, we know

\begin{align}
\hat{\beta} = \frac{\textrm{cov}(x,y)}{\sigma_x^2},
\end{align}

and

\begin{align}
r = \frac{\textrm{cov}(x,y)}{\sigma_x \sigma_y}.
\end{align}

Thus

\begin{align}
R^2 &= r^2, \\
&= \left(\frac{\hat{\beta}\sigma_x^2}{\sigma_x \sigma_y}\right)^2, \\
&= \left(\hat{\beta}\frac{\sigma_x}{\sigma_y}\right)^2.
\end{align}

## Subsets

```python
[rna_meta.in_eqtl.sum()], [qvalues.perm_sig.sum()]
```

    ([215], [5816])

```python
cols = [40, 50, 60, 70, 80, 90, 100]
pvals = {}
for i in cols:
    fn = os.path.join(ciepy.root, 'output', 'eqtl_processing',
                      'unrelated_eqtls_{}'.format(i), 'min_pvalues_corrected.tsv')
    pvals[i] = pd.read_table(fn, index_col=0).min_pval_bf
pvals = pd.DataFrame(pvals)

(pvals < 0.05).sum().plot()
plt.scatter(pvals.columns, (pvals < 0.05).sum().values)
#plt.scatter([rna_meta.in_eqtl.sum()], [qvalues.perm_sig.sum()])
plt.xlim(35, 105)
plt.xlabel('Number of samples')
plt.ylabel('Number of eGenes');
```

## pwr

TODO: Below, I calculate the genotype stdev by pulling the genotypes out of the VCF files.
However, I could more easily just read the EMMAX results files and use the `genocnt` column.
This would save time for this step and the power step since I could avoid calculating power
for some of the GTEx variants I don't have. I also need to make sure to calculate genotype stdev
for each set of subjects (40, 50, 60, etc.).

```python
def pwr_df(df, u, v):
    """
    Execute pwr.f2.test. The dataframe df must have a column f2.
    """
    import rpy2.robjects as ro
    ro.r('suppressPackageStartupMessages(library(pwr))')
    ro.globalenv['df'] = df
    ro.globalenv['u'] = u
    ro.globalenv['v'] = v
    ro.r('p = rep(1, dim(df)[2])')
    ro.r('for(i in 1:dim(df)[1]) {'
         'p[i] = pwr.f2.test(u=u, v=v, f2=df[i, "f2"], '
         'sig.level=0.001, power=NULL)["power"]}')
    p = ro.globalenv['p']
    power = [x[0] for x in p]
    return power
```

```python
def pwr(f2, u, v):
    """
    Execute pwr.f2.test.
    """
    import rpy2.robjects as ro
    ro.r('suppressPackageStartupMessages(library(pwr))')
    ro.globalenv['f2'] = f2
    ro.globalenv['u'] = u
    ro.globalenv['v'] = v
    ro.r('p = pwr.f2.test(u=u, v=v, f2=f2, sig.level=0.001, power=NULL)["power"]')
    p = list(ro.globalenv['p'])[0][0]
    return p
```
def calc_stdev(chrom, pos):
    res = vcf_reader.fetch(chrom, pos, pos)
    try:
        r = res.next()
        out = np.std([0] * r.num_hom_ref + [1] * r.num_het + [2] * r.num_hom_alt)
    except:
        out = np.nan
    return '{}:{}'.format(chrom, int(pos)), out
Here's an example of using `pwr.f2.test`.

```r
%%R

suppressPackageStartupMessages(library(pwr))
pwr.f2.test(u=2, v=117, f2=0.060995, sig.level=0.001, power=NULL)
```

         Multiple regression power calculation

                  u = 2
                  v = 117
                 f2 = 0.060995
          sig.level = 0.001
              power = 0.1766086

### Genotype stdev

```python
def calc_stdev(x):
    hr, het, ha = [int(x) for x in x.split('/')]
    out = np.std([0] * hr + [1] * het + [2] * ha)
    return out
```

```python
stdev = {}
for i in np.arange(40, 110, 10):
    k = 'geno_stdev_{}'.format(i)
    out = os.path.join(outdir, '{}.tsv'.format(k))
    if not os.path.exists(out):
        fn = os.path.join(ciepy.root, 'output', 'eqtl_processing',
                          'unrelated_eqtls_{}'.format(i), 'top_snv_results_sorted.tsv.gz')
        df = pd.read_table(fn, header=None, low_memory=False)
        se = df[7].apply(lambda x: calc_stdev(x))
        se.index = 'chr' + df[0].astype(str) + ':' + df[1].astype(str)
        stdev[k] = se
        stdev[k].to_csv(out, sep='\t')
    else:
        stdev[k] = pd.read_table(out, index_col=0, header=None, squeeze=True)
```

```python
fn = os.path.join(outdir, 'genotype_stdevs.tsv')
if not os.path.exists(fn):
    ind = []
    for v in stdev.values():
        ind += list(v.index)
    ind = set(ind)
    for k in stdev:
        stdev[k] = stdev[k][ind]
    geno_stdev = pd.DataFrame(stdev)
    geno_stdev.to_csv(fn, sep='\t')
else:
    geno_stdev = pd.read_table(fn, index_col=0)
```

### Geuvadis
url = ('ftp://ftp.ebi.ac.uk/pub/databases/microarray/data/experiment/'
       'GEUV/E-GEUV-1/analysis_results/EUR373.gene.cis.FDR5.all.rs137.txt.gz')
fn = os.path.join(private_outdir, 'EUR373.gene.cis.FDR5.all.rs137.txt')
if not os.path.exists(fn):
    !curl {url} > {fn + '.gz'}
    !gunzip {fn + '.gz'}

geuvadis_fn = os.path.join(outdir, 'geuvadis_stdev.tsv')
if not os.path.exists(geuvadis_fn):
    geuvadis = pd.read_table(os.path.join(private_outdir, 'EUR373.gene.cis.FDR5.all.rs137.txt'),
                             low_memory=False)
    # Remove genes that I didn't test.
    geuvadis['ens_gene_id'] = geuvadis.GENE_ID.apply(lambda x: x.split('.')[0])
    tested = pd.DataFrame(True, index=[x.split('.')[0] for x in exp.index], columns=['tested'])
    geuvadis = geuvadis.merge(tested, left_on='ens_gene_id', right_index=True)
    # Remove indels.
    geuvadis = geuvadis[geuvadis.SNPpos - geuvadis.SNPpos.astype(int) == 0]
    # Add standard deviations
    geuvadis['chr_pos'] = ('chr' + geuvadis.CHR_SNP.astype(str) + ':' +
                           geuvadis.SNPpos.astype(int).astype(str))
    geuvadis = geuvadis.merge(geno_stdev, left_on='chr_pos', right_index=True)
    # Add R^2 and f^2 values
    # TODO: I need to calculate R^2 using my genotype stdevs but since they
    # don't provide beta values I can't do that.
### GTEx

The power calculations are pretty slow if we calculate for each individual
association. To speed this up, we'll precalculate the power values for
different $f2$ values and figure out for which values the power is one.

```python
vals = []
for i in np.arange(0, 10, 0.01):
    vals.append(pwr(i, 2, 37))
plt.scatter(np.arange(0, 10, 0.01), vals)
plt.xlabel('$f2$')
plt.ylabel('Power');
```

We can see that even for 40 samples the power goes to one near $f^2 = 1$. This
will happen even faster for more samples.

```python
vals = []
for i in np.arange(0, 2, 0.001):
    vals.append(pwr(i, 2, 37))
plt.scatter(np.arange(0, 2, 0.001), vals)
plt.xlabel('$f2$')
plt.ylabel('Power');
```

I'll go in increments of $0.001$ until I reach power of one.

```python
def precalculate_power(u, v):
    i = 0
    cur_power = 0
    ind = []
    vals = []
    while cur_power < 1.:
        cur_power = pwr(i, u, v)
        vals.append(cur_power)
        i += 0.001
    pdf = pd.DataFrame(vals, index=np.arange(0, i, 0.001), columns=['power'])
    return pdf
```

```python
# Precalculate power values.
power = {}
for i in np.arange(40, 110, 10):
    power[i] = precalculate_power(2, i - 3)
    power[i].columns = ['power_{}'.format(i)]
```

```python
def get_gene_res(gene, dy, gtex_ind):
    fn = os.path.join(dy, gene, '{}.tsv'.format(gene))
    res = ciepy.read_emmax_output(fn)
    res.index = 'chr' + res.CHROM.astype(str) + ':' + res.BEG.astype(str) + ':' + gene
    if type(gtex_ind) == str:
        ind = [gtex_ind]
    else:
        ind = gtex_ind.values
    ind = set(res.index) & set(ind)
    res = res.ix[ind]
    return res.PVALUE < 0.001
```

The code below takes a while to run. The longest part is reading
through all of my results for the individual genes for each eQTL
analysis. This could probably be sped up by running on an IPython
cluster.

```python
fns = glob.glob('/publicdata/gtex_v6/*.snpgenes')
out_fns = glob.glob(os.path.join(outdir, 'gtex_*_power.tsv'))
if len(out_fns) != fns:
    gtex_results = {}
    gtex_sig = []
    for fn in fns:
        k = os.path.split(fn)[1][:-18]
        t = pd.read_table(fn, index_col=0, low_memory=False)
        ref = [x.split('_')[2] for x in t.index]
        alt = [x.split('_')[2] for x in t.index]
        t['ref_length'] = [len(x) for x in ref]
        t['alt_length'] = [len(x) for x in alt]
        t = t[t.ref_length == 1]
        t = t[t.alt_length == 1]
        t['chr_pos'] = ('chr' + t.snp_chrom.astype(str) + ':' + t.snp_pos.astype(str))
        t = t.merge(geno_stdev, left_on='chr_pos', right_index=True)
        for i in np.arange(40, 110, 10):
            t['R2_{}'.format(i)] = (t.beta * t['geno_stdev_{}'.format(i)]) ** 2
            # We'll round the f2 values so I can use my precalculated power estimates.
            t['f2_{}'.format(i)] = np.round(t['R2_{}'.format(i)] / (1. - t['R2_{}'.format(i)]), 3)
            # f2 must be positive. Negative is probably a numeric problem.
            t = t[t['f2_{}'.format(i)] > 0]
            cpower = power[i].ix[set(t['f2_{}'.format(i)])].dropna()
            t = t.merge(cpower, left_on='f2_{}'.format(i), right_index=True, how='outer')
            t.ix[t['power_{}'.format(i)].isnull(), 'power_{}'.format(i)] = 1
        tt = 'chr' + t.snp_chrom.astype(str) + ':' + t.snp_pos.astype(str) + ':' + t.gene
        tt.index = t.gene
        gtex_sig.append(tt)
        gtex_results[k] = t
    gtex_sig = pd.concat(gtex_sig).drop_duplicates()
    gtex_sig = gtex_sig.ix[set(gtex_sig.index) & set(qvalues.index)]

    my_gtex_results = pd.DataFrame(index=gtex_sig.values)

#     for i in np.arange(40, 110, 10):
#         col = 'sig_{}'.format(i)
#         my_gtex_results[col] = np.nan
#         dy = os.path.join(ciepy.root, 'private_output', 'run_eqtl_analysis',
#                           'unrelated_eqtls_{}'.format(i), 'gene_results')
#         for gene in set(gtex_sig.index):
#             fn = os.path.join(dy, gene, '{}.tsv'.format(gene))
#             res = ciepy.read_emmax_output(fn)
#             res.index = 'chr' + res.CHROM.astype(str) + ':' + res.BEG.astype(str) + ':' + gene
#             t = gtex_sig.ix[gene]
#             if type(t) == str:
#                 ind = [t]
#             else:
#                 ind = t.values
#             ind = set(res.index) & set(ind)
#             res = res.ix[ind]
#             my_gtex_results.ix[ind, col] = (res.PVALUE < 0.001).values
    from ipyparallel import Client
    parallel_client = Client(profile='parallel')
    dview = parallel_client[:]
    print('Cluster has {} engines.'.format(len(parallel_client.ids)))

    with dview.sync_imports():
        import numpy
        import ciepy
        import os
        import vcf

    %px np = numpy

    dview.push(dict(get_gene_res=get_gene_res))

    todo = []
    for g in set(gtex_sig.index):
        todo.append([g, gtex_sig[g]])

    my_gtex_results = {}
    for i in np.arange(40, 110, 10):
        dy = os.path.join(ciepy.root, 'private_output', 'run_eqtl_analysis',
                          'unrelated_eqtls_{}'.format(i), 'gene_results')
        dview.push(dict(dy=dy))
        t = dview.map_sync(lambda x: get_gene_res(x[0], dy, x[1]), todo)
        tt = pd.concat(t)
        my_gtex_results['sig_{}'.format(i)] = tt
    my_gtex_results = pd.DataFrame(my_gtex_results)

    for k in gtex_results.keys():
        t = gtex_results[k]
        t['chr_pos_gene'] = t.chr_pos + ':' + t.gene
        t = t.merge(my_gtex_results, left_on='chr_pos_gene', right_index=True)
        for i in np.arange(40, 110, 10):
            t['power_{}_bin'.format(i)] = pd.cut(t['power_{}'.format(i)], np.arange(0, 1.01, 0.01)).values
        gtex_results[k] = t
        out = os.path.join(outdir, 'gtex_{}_power.tsv'.format(k))
        t.to_csv(out, sep='\t')
```

```python
fig,axs = plt.subplots(len(gtex_results) / 2, 2, figsize=(8, 40))
for i,k in enumerate(gtex_results.keys()):
    ax = axs[i / 2, i % 2]
    t = gtex_results[k]
    t = t.sort_values(by='p_value')
    t = t.drop_duplicates('gene')
    t.R2_100.hist(bins=np.arange(0, 1.01, 0.01), ax=ax)
    ax.set_title(k);
    ax.set_ylabel('Number of eGenes')
    ax.set_xlabel('$R^2$');
plt.tight_layout()
```

```python
def power_plots(tdf):
    for i in np.arange(40, 110, 10):
        df = pd.crosstab(tdf['power_{}_bin'.format(i)], tdf['sig_{}'.format(i)])
        df['percent'] = df[True] / df.sum(axis=1)
        plt.figure()
        df.percent.plot()
        plt.title(i)
        plt.ylabel('Observed power')
        plt.xlabel('Estimated power');
```

```python
import statsmodels.api as sm
```

```python
t = gtex_results['Thyroid']
i = 40
```

```python
df = pd.crosstab(t['power_{}_bin'.format(i)], t['sig_{}'.format(i)])
df['percent'] = df[True] / df.sum(axis=1)
```

```python
def power_plots(tdf):
    with sns.plotting_context('notebook', font_scale=1.5):
        plt.figure(figsize=(8, 8))
        for n,i in enumerate(np.arange(40, 110, 10)):
            df = pd.crosstab(tdf['power_{}_bin'.format(i)], tdf['sig_{}'.format(i)])
            df['percent'] = df[True] / df.sum(axis=1)
            xs = np.arange(0, 1, 0.01)
            smoothed = pd.DataFrame(sm.nonparametric.lowess(df.percent, xs),
                                    columns=['x', i])
            smoothed.index = smoothed['x']

            smoothed[i].plot(color=cpb.analysis.tableau20[2 * n])
            plt.title(i)
            plt.ylabel('Observed power')
            plt.xlabel('Estimated power')
        plt.title('')
        plt.legend(loc='upper left');
```

```python
power_plots(gtex_results['Liver'])
```

```python
power_plots(gtex_results['Thyroid'])
```

```python
power_plots(gtex_results['Lung'])
```

```python
t = gtex_results['Lung']
```

```python
t.power_40.hist()
```

```python
t.power_100.hist()
```

```python
prop_p100 = []
for k in gtex_results.keys():
    t = gtex_results[k]
    t = t[t.power_100 == 1]
    prop_p100.append(t.sig_100.sum() / float(t.shape[0]))
prop_p100 = pd.Series(prop_p100, index=gtex_results.keys())
prop_p100.sort_values(inplace=True)
```

```python
prop_p100.plot(kind='barh')
plt.xlabel('Proportion of genes with expected power == 1 detected')
plt.tight_layout()
```

This is still affected by sample size since tissues with small
sample size have few eQTLs with large effect sizes.


# Knowledge Distillation

В этом ноутбуке содержится описание моих основных экспериментов по knowledge distillation. Здесь будут в основном графики и описания экспериментов, весь код для обучения и оценки качества моделей вынесен в отдельные модули. Проще всего запустить в Google Colab, тогда не нужно дополнительно устанавливать никаких зависимостей. Если хотите запустить решение на своей локальной машине, то следуйте инструкциям в README.md.

Выполните следующие две ячейки, если работаете в Google Colab:

```python
! git clone https://github.com/stdereka/knowledge-distillation.git
! cp -r knowledge-distillation/dark_knowledge/ .
! cp -r knowledge-distillation/experiments/ .
```

    Cloning into 'knowledge-distillation'...
    remote: Enumerating objects: 257, done.
    remote: Counting objects: 100% (257/257), done.
    remote: Compressing objects: 100% (205/205), done.
    remote: Total 257 (delta 125), reused 156 (delta 46), pack-reused 0
    Receiving objects: 100% (257/257), 4.12 MiB | 26.20 MiB/s, done.
    Resolving deltas: 100% (125/125), done.

```python
import sys
sys.path.insert(1, "/content/knowledge-distillation")
```

## 1. Описание задачи

Метод KD был изначально предложен в [статье](https://arxiv.org/pdf/1503.02531.pdf). Для задачи многоклассовой классификации утверждается, что вероятности классов, предсказанные моделью-учителем, могут быть исползованы для обучения более простой модели-студента с меньшим числом параметров путём модификации лосса. Для многоклассовой классификации предлагается следующая функция потерь (здесь индекс $i$ пробегает по всем классам и по всем объектам):

\begin{equation}
E=-T^{2} (1 - \alpha) \sum_{i} \hat{y}_{i}(\mathbf{x} \mid T) \log y_{i}(\mathbf{x} \mid T)- \alpha\sum_{i} \bar{y}_{i} \log y_{i}(\mathbf{x} \mid 1)
\end{equation}

где $y_{i}(\mathbf{x} \mid T)=\frac{e^{\frac{z_{i}(\mathbf{x})}{T}}}{\sum_{j} e^{\frac{z_{j}(\mathrm{x})}{T}}}$ - предсказания модели при температуре $T$ (она требуется для того, чтобы предсказания сильнее отличались от 0 и 1 и несли больше информации), $\hat{y}_{i}(\mathbf{x} \mid T)$ - "мягкие" метки учителя, $\bar{y}_{i}$ - "жёсткие" groundtruth метки. $\alpha$ - весовой параметр для регулирования вклада жёсткой и мягкой кроссэнтропий в градиент. Авторы статьи пишут, что всё должно хорошо работать при $\alpha$ близком к нулю.

## 2. Датасет Imagewoof

Imagewoof - датасет из изображений собак, относящимся к разным породам. На нём можно решать задачу многоклассовой классификации. Я выбрал его для экспериментов по нескольким причинам. Во-первых, классы трудноразделимы, слишком простые модели не смогут эффективно решить задачи, поэтому есть шансы пронаблюдать эффекты KD на более сложных моделях. Во-вторых, по Imagewoof не так много публикаций на тему KD, это значит, что результаты экспериментов мне заранее не известны.

Сначала я хотел начать с воспроизведения результатов статьи на датасете MNIST, но это уже сделали за меня, например, [тут](https://github.com/peterliht/knowledge-distillation-pytorch). Я решил сразу начать со своих экспериментов на более сложной задаче.

## 3. Общая схема эксперимента

1. **Выбрать и обучить модель-учитель.** Эта модель должна содержать большое число параметров (возможно, даже избыточное) и быть сильно регуляризована (в задачах CV подойдут сильные дропауты и аугментации), а также выдавать очень хорошее (по меркам задачи) качество на отложенной выборке.

2. **Обучение и оценка качества модели-студента с Distillation Loss.** Модель-студент содержит ощутимо меньше параметров, чем модель-учитель и никак не регуляризована. На этом этапе можно играть с параметрами $\alpha$, $T$, размером обучающей выборки и много чем ещё.

3. **Обучить модель-студента без помощи учителя.** Здесь очень важно, чтобы процесс обучения отличался только лоссом, это позволит исключить вклад всех прочих факторов.

4. **Сравнить метрики в пунктах 2 и 3**, построить выводы об эффективности KD в текущем эксперименте.

Из бесплатных ресурсов в моём распоряженни только Google Colab, это очень неудобно с точки зрения разработки и воспроизводимости результатов, но я смог найти выход и прийти к следующей схеме. Код пишется и тестируется локально на моём компьютере. Когда всё готово, ноутбук с параметрами отправляется на Google Сolab, подтягивает код с GitHub и выполняет эксперимент. Затем код, результаты и параметры сохраняются в одном коммите, это позволяет при необходимости вернуться и воспроизвести эксперимент.

## 4. Демонстрационный эксперимент

В этой части я проведу эксперимент согласно пункту 3.

```python
# Load data
! wget https://s3.amazonaws.com/fast-ai-imageclas/imagewoof2-320.tgz
! tar zxf imagewoof2-320.tgz
# Check GPU
! nvidia-smi
```

    --2020-09-23 16:49:23--  https://s3.amazonaws.com/fast-ai-imageclas/imagewoof2-320.tgz
    Resolving s3.amazonaws.com (s3.amazonaws.com)... 52.216.142.214
    Connecting to s3.amazonaws.com (s3.amazonaws.com)|52.216.142.214|:443... connected.
    HTTP request sent, awaiting response... 200 OK
    Length: 328288506 (313M) [application/x-tar]
    Saving to: ‘imagewoof2-320.tgz’

    imagewoof2-320.tgz  100%[===================>] 313.08M  20.6MB/s    in 12s

    2020-09-23 16:49:36 (26.7 MB/s) - ‘imagewoof2-320.tgz’ saved [328288506/328288506]

    Wed Sep 23 16:49:39 2020
    +-----------------------------------------------------------------------------+
    | NVIDIA-SMI 450.66       Driver Version: 418.67       CUDA Version: 10.1     |
    |-------------------------------+----------------------+----------------------+
    | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
    | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
    |                               |                      |               MIG M. |
    |===============================+======================+======================|
    |   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
    | N/A   44C    P8     9W /  70W |      0MiB / 15079MiB |      0%      Default |
    |                               |                      |                 ERR! |
    +-------------------------------+----------------------+----------------------+

    +-----------------------------------------------------------------------------+
    | Processes:                                                                  |
    |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
    |        ID   ID                                                   Usage      |
    |=============================================================================|
    |  No running processes found                                                 |
    +-----------------------------------------------------------------------------+

```python
import torch
import random
import numpy as np
from pathlib import Path
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import accuracy_score
from torch.utils.data import DataLoader
from torch import nn
import matplotlib.pyplot as plt
import os
from training import train, predict, DistillationLoss
from models import *
from datasets import ImageDataset
from itertools import product
```

```python
"""
Define globals and seed whatever I can seed
"""
def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

train_on_gpu = torch.cuda.is_available()
if not train_on_gpu:
    print('Training on CPU')
    DEVICE = torch.device("cpu")
else:
    print('Training on GPU')
    DEVICE = torch.device("cuda")

SEED = 0
seed_everything(SEED)

TRAIN_DIR = Path('./imagewoof2-320/train')
TEST_DIR = Path('./imagewoof2-320/val')

train_val_files = sorted(list(TRAIN_DIR.rglob('*.JPEG')))
test_files = sorted(list(TEST_DIR.rglob('*.JPEG')))

train_val_labels = [path.parent.name for path in train_val_files]
LABEL_ENCODER = LabelEncoder()
LABEL_ENCODER.fit(train_val_labels)

N_CLASSES = LABEL_ENCODER.classes_.shape[0]
```

    Training on GPU

```python
def plot_training_history(history):
    fig, ax = plt.subplots(1, 2, figsize=(20, 5))
    ax = ax.flatten()

    ax[0].plot(history[:, 0], label="Train")
    ax[0].plot(history[:, 2], label="Val")
    ax[0].set_xlabel("Epoch")
    ax[0].set_ylabel("Loss")
    ax[0].legend()
    ax[0].grid()

    ax[1].plot(history[:, 1], label="Train")
    ax[1].plot(history[:, 3], label="Val")
    ax[1].set_xlabel("Epoch")
    ax[1].set_ylabel("Accuracy")
    ax[1].legend()
    ax[1].grid()

    plt.show()
```

### 4.1. Выбор и обучение модели-учителя

Учитель - Resnet101. Я добавил два скрытых полносвязных слоя по 2048 юнитов. Такая модель явно избыточна, очень быстро переобучается и требует адский дропаут ($p=0.95$), чтобы не переобучиться и показать хорошие результаты на валидационной выборке. При обучении этой модели я также задействовал сильные аугментации.

Для дальнейших экспериментов (например, с размером обучающей выборки для модели-студента) могут понадобиться предсказания учителя на всей обучающей выборке. Чтобы не допустить утечки данных я использую 1x4 кроссвалидацию со стратификацией по классам, итоговые предсказания - это out-of-fold предсказания моделей-учителей, соответствующих разным фолдам. Таким образом, для получения предсказаний на всём датасете нужно обучить учителя целых 4 раза, зато после этого можно надолго о нём забыть и заняться непосредственно KD.

```python
N_SPLITS = 4

seed_everything(SEED)

cv = StratifiedKFold(n_splits=N_SPLITS, shuffle=True, random_state=SEED)
ground = np.array(ImageDataset(train_val_files, LABEL_ENCODER).labels)

# Array for storing OOF logits
oof_logits = np.empty((len(train_val_files), N_CLASSES), float)

labels = LABEL_ENCODER.transform(train_val_labels)

test_dataset = ImageDataset(test_files, LABEL_ENCODER)
test_loader = DataLoader(test_dataset, shuffle=False, batch_size=64)

# Array for storing test dataset logits
test_logits = np.zeros((len(test_dataset), N_CLASSES), float)

for fold, (train_idx, val_idx) in enumerate(cv.split(train_val_files, labels)):
    print('Training on fold', fold + 1)

    val_dataset = ImageDataset(np.array(train_val_files)[val_idx], LABEL_ENCODER)
    train_dataset = ImageDataset(np.array(train_val_files)[train_idx], LABEL_ENCODER, augs=True)

    model = resnet101_teacher(N_CLASSES, DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=0.0003)
    criterion = nn.CrossEntropyLoss()
    history = train(train_dataset, val_dataset, model=model, epochs=20,
                              batch_size=64, device=DEVICE, opt=opt, criterion=criterion)

    val_loader = DataLoader(val_dataset, shuffle=False, batch_size=64)
    labels_val = ground[val_idx]

    logits_val = predict(model, val_loader, DEVICE, logit=True)
    oof_logits[val_idx] = logits_val
    preds_val = np.argmax(logits_val, axis=1)

    # Prediction on test set is an average of N_SPLITS models
    test_logits += predict(model, test_loader, DEVICE, logit=True)/N_SPLITS

    print(f'Fold {fold + 1} accuracy score:', accuracy_score(labels_val, preds_val))

oof_preds = np.argmax(oof_logits, axis=1)
print('OOF accuracy score:', accuracy_score(ground, oof_preds))

# Save results
os.makedirs("./dark_knowledge", exist_ok=True)
np.save("./dark_knowledge/resnet101_train_imagewoof.npy", oof_logits)
np.save("./dark_knowledge/resnet101_test_imagewoof.npy", test_logits)
```

    Training on fold 1

    Downloading: "https://download.pytorch.org/models/resnet101-5d3b4d8f.pth" to /root/.cache/torch/hub/checkpoints/resnet101-5d3b4d8f.pth

    Epoch: 1, train_loss: 2.7788, val_loss: 2.2557, train_acc: 0.1169, val_acc: 0.4945
    Epoch: 2, train_loss: 1.4853, val_loss: 0.3969, train_acc: 0.8097, val_acc: 0.9278
    Epoch: 3, train_loss: 0.3147, val_loss: 0.1817, train_acc: 0.9170, val_acc: 0.9508
    Epoch: 4, train_loss: 0.2246, val_loss: 0.1467, train_acc: 0.9279, val_acc: 0.9526
    Epoch: 5, train_loss: 0.2063, val_loss: 0.1502, train_acc: 0.9334, val_acc: 0.9504
    Epoch: 6, train_loss: 0.1927, val_loss: 0.1605, train_acc: 0.9362, val_acc: 0.9504
    Epoch: 7, train_loss: 0.1982, val_loss: 0.1504, train_acc: 0.9341, val_acc: 0.9499
    Epoch: 8, train_loss: 0.1852, val_loss: 0.1447, train_acc: 0.9387, val_acc: 0.9521
    Epoch: 9, train_loss: 0.1910, val_loss: 0.1451, train_acc: 0.9400, val_acc: 0.9539
    Epoch: 10, train_loss: 0.1836, val_loss: 0.1380, train_acc: 0.9387, val_acc: 0.9544
    Epoch: 11, train_loss: 0.1727, val_loss: 0.1329, train_acc: 0.9402, val_acc: 0.9601
    Epoch: 12, train_loss: 0.1752, val_loss: 0.1386, train_acc: 0.9430, val_acc: 0.9530
    Epoch: 13, train_loss: 0.1765, val_loss: 0.1256, train_acc: 0.9406, val_acc: 0.9601
    Epoch: 14, train_loss: 0.1666, val_loss: 0.1510, train_acc: 0.9464, val_acc: 0.9477
    Epoch: 15, train_loss: 0.1795, val_loss: 0.1458, train_acc: 0.9382, val_acc: 0.9517
    Epoch: 16, train_loss: 0.1622, val_loss: 0.1456, train_acc: 0.9431, val_acc: 0.9495
    Epoch: 17, train_loss: 0.1627, val_loss: 0.1430, train_acc: 0.9437, val_acc: 0.9513
    Epoch: 18, train_loss: 0.1541, val_loss: 0.1360, train_acc: 0.9468, val_acc: 0.9553
    Epoch: 19, train_loss: 0.1705, val_loss: 0.1337, train_acc: 0.9413, val_acc: 0.9521
    Epoch: 20, train_loss: 0.1613, val_loss: 0.1361, train_acc: 0.9439, val_acc: 0.9539
    Fold 1 accuracy score: 0.9539211342490032
    Training on fold 2
    Epoch: 1, train_loss: 2.8263, val_loss: 2.2559, train_acc: 0.1129, val_acc: 0.5421
    Epoch: 2, train_loss: 1.5103, val_loss: 0.3935, train_acc: 0.7619, val_acc: 0.9260
    Epoch: 3, train_loss: 0.3205, val_loss: 0.1710, train_acc: 0.9078, val_acc: 0.9495
    Epoch: 4, train_loss: 0.2329, val_loss: 0.1350, train_acc: 0.9264, val_acc: 0.9592
    Epoch: 5, train_loss: 0.2046, val_loss: 0.1244, train_acc: 0.9338, val_acc: 0.9641
    Epoch: 6, train_loss: 0.1978, val_loss: 0.1242, train_acc: 0.9350, val_acc: 0.9597
    Epoch: 7, train_loss: 0.1974, val_loss: 0.1042, train_acc: 0.9317, val_acc: 0.9654
    Epoch: 8, train_loss: 0.1886, val_loss: 0.1045, train_acc: 0.9362, val_acc: 0.9681
    Epoch: 9, train_loss: 0.2034, val_loss: 0.1297, train_acc: 0.9325, val_acc: 0.9574
    Epoch: 10, train_loss: 0.1878, val_loss: 0.1153, train_acc: 0.9365, val_acc: 0.9605
    Epoch: 11, train_loss: 0.2074, val_loss: 0.1181, train_acc: 0.9340, val_acc: 0.9605
    Epoch: 12, train_loss: 0.1742, val_loss: 0.1274, train_acc: 0.9391, val_acc: 0.9535
    Epoch: 13, train_loss: 0.1793, val_loss: 0.1222, train_acc: 0.9388, val_acc: 0.9597
    Epoch: 14, train_loss: 0.1784, val_loss: 0.1198, train_acc: 0.9391, val_acc: 0.9605
    Epoch: 15, train_loss: 0.1824, val_loss: 0.1066, train_acc: 0.9400, val_acc: 0.9650
    Epoch: 16, train_loss: 0.1666, val_loss: 0.1093, train_acc: 0.9447, val_acc: 0.9614
    Epoch: 17, train_loss: 0.1745, val_loss: 0.1219, train_acc: 0.9391, val_acc: 0.9623
    Epoch: 18, train_loss: 0.1702, val_loss: 0.1145, train_acc: 0.9443, val_acc: 0.9650
    Epoch    19: reducing learning rate of group 0 to 9.0000e-05.
    Epoch: 19, train_loss: 0.1643, val_loss: 0.1224, train_acc: 0.9421, val_acc: 0.9619
    Epoch: 20, train_loss: 0.1441, val_loss: 0.1026, train_acc: 0.9490, val_acc: 0.9650
    Fold 2 accuracy score: 0.9649822695035462
    Training on fold 3
    Epoch: 1, train_loss: 2.7962, val_loss: 2.2467, train_acc: 0.1194, val_acc: 0.3240
    Epoch: 2, train_loss: 1.4524, val_loss: 0.3583, train_acc: 0.7653, val_acc: 0.9486
    Epoch: 3, train_loss: 0.3053, val_loss: 0.1763, train_acc: 0.9153, val_acc: 0.9512
    Epoch: 4, train_loss: 0.2319, val_loss: 0.1384, train_acc: 0.9248, val_acc: 0.9579
    Epoch: 5, train_loss: 0.2234, val_loss: 0.1533, train_acc: 0.9281, val_acc: 0.9512
    Epoch: 6, train_loss: 0.2122, val_loss: 0.1346, train_acc: 0.9307, val_acc: 0.9574
    Epoch: 7, train_loss: 0.1964, val_loss: 0.1336, train_acc: 0.9328, val_acc: 0.9601
    Epoch: 8, train_loss: 0.1927, val_loss: 0.1450, train_acc: 0.9332, val_acc: 0.9539
    Epoch: 9, train_loss: 0.1935, val_loss: 0.1457, train_acc: 0.9341, val_acc: 0.9543
    Epoch: 10, train_loss: 0.1887, val_loss: 0.1670, train_acc: 0.9341, val_acc: 0.9508
    Epoch: 11, train_loss: 0.1781, val_loss: 0.1352, train_acc: 0.9409, val_acc: 0.9552
    Epoch: 12, train_loss: 0.1689, val_loss: 0.1508, train_acc: 0.9409, val_acc: 0.9512
    Epoch: 13, train_loss: 0.1779, val_loss: 0.1306, train_acc: 0.9390, val_acc: 0.9588
    Epoch: 14, train_loss: 0.1757, val_loss: 0.1378, train_acc: 0.9437, val_acc: 0.9583
    Epoch: 15, train_loss: 0.1649, val_loss: 0.1489, train_acc: 0.9428, val_acc: 0.9552
    Epoch: 16, train_loss: 0.1726, val_loss: 0.1308, train_acc: 0.9422, val_acc: 0.9570
    Epoch: 17, train_loss: 0.1744, val_loss: 0.1609, train_acc: 0.9414, val_acc: 0.9486
    Epoch    18: reducing learning rate of group 0 to 9.0000e-05.
    Epoch: 18, train_loss: 0.1671, val_loss: 0.1264, train_acc: 0.9447, val_acc: 0.9592
    Epoch: 19, train_loss: 0.1432, val_loss: 0.1227, train_acc: 0.9533, val_acc: 0.9588
    Epoch: 20, train_loss: 0.1336, val_loss: 0.1258, train_acc: 0.9530, val_acc: 0.9628
    Fold 3 accuracy score: 0.9627659574468085
    Training on fold 4
    Epoch: 1, train_loss: 2.7867, val_loss: 2.2422, train_acc: 0.1158, val_acc: 0.7057
    Epoch: 2, train_loss: 1.4403, val_loss: 0.3659, train_acc: 0.8127, val_acc: 0.9211
    Epoch: 3, train_loss: 0.3105, val_loss: 0.2038, train_acc: 0.9170, val_acc: 0.9326
    Epoch: 4, train_loss: 0.2357, val_loss: 0.1487, train_acc: 0.9239, val_acc: 0.9552
    Epoch: 5, train_loss: 0.2257, val_loss: 0.1388, train_acc: 0.9239, val_acc: 0.9539
    Epoch: 6, train_loss: 0.2045, val_loss: 0.1317, train_acc: 0.9326, val_acc: 0.9588
    Epoch: 7, train_loss: 0.2040, val_loss: 0.1624, train_acc: 0.9294, val_acc: 0.9499
    Epoch: 8, train_loss: 0.1919, val_loss: 0.1336, train_acc: 0.9393, val_acc: 0.9570
    Epoch: 9, train_loss: 0.1900, val_loss: 0.1292, train_acc: 0.9390, val_acc: 0.9628
    Epoch: 10, train_loss: 0.1891, val_loss: 0.1332, train_acc: 0.9354, val_acc: 0.9579
    Epoch: 11, train_loss: 0.1874, val_loss: 0.1245, train_acc: 0.9350, val_acc: 0.9610
    Epoch: 12, train_loss: 0.1832, val_loss: 0.1162, train_acc: 0.9385, val_acc: 0.9654
    Epoch: 13, train_loss: 0.1843, val_loss: 0.1115, train_acc: 0.9338, val_acc: 0.9654
    Epoch: 14, train_loss: 0.1734, val_loss: 0.1278, train_acc: 0.9403, val_acc: 0.9628
    Epoch: 15, train_loss: 0.1782, val_loss: 0.1357, train_acc: 0.9381, val_acc: 0.9552
    Epoch: 16, train_loss: 0.1641, val_loss: 0.1310, train_acc: 0.9450, val_acc: 0.9588
    Epoch: 17, train_loss: 0.1654, val_loss: 0.1181, train_acc: 0.9458, val_acc: 0.9623
    Epoch: 18, train_loss: 0.1640, val_loss: 0.1267, train_acc: 0.9409, val_acc: 0.9592
    Epoch: 19, train_loss: 0.1614, val_loss: 0.1378, train_acc: 0.9434, val_acc: 0.9597
    Epoch: 20, train_loss: 0.1660, val_loss: 0.1302, train_acc: 0.9403, val_acc: 0.9574
    Fold 4 accuracy score: 0.9574468085106383
    OOF accuracy score: 0.9597783933518006

### 4.2. Модель-студент. Дистилляция знаний

Студент - Resnet18 с двумя дополнительными скрытыми полносвязными слоями по 128 юнитов. При обучении этой модели я не задействую ни дропауты, ни аугментации. Спойлер: это необходимо, потому что эффект от KD в этой задаче очень маленький, и чтобы его обнаружить, нужно исключить все прочие факторы, влияющие на качество модели.

```python
os.makedirs("./experiments/hyperparams", exist_ok=True)
seed_everything(SEED)

resnet18 = resnet18_student2(N_CLASSES, DEVICE)
trainable = get_number_of_params(resnet18, trainable=True)
total = get_number_of_params(resnet18, trainable=False)
print(f"{trainable} (of {total}) trainable params")
train_dataset = ImageDataset(train_val_files, LABEL_ENCODER, teacher_labels="./dark_knowledge/resnet101_train_imagewoof.npy")
test_dataset = ImageDataset(test_files, LABEL_ENCODER, teacher_labels="./dark_knowledge/resnet101_test_imagewoof.npy")
opt = torch.optim.Adam(resnet18.parameters(), lr=0.0001)
criterion = DistillationLoss(alpha=0.1, temperature=4.0)
history_resnet18 = train(train_dataset, test_dataset, resnet18, 50, 64, DEVICE, opt, criterion)
plot_training_history(history_resnet18)
np.save("./experiments/hyperparams/history_resnet18_T_4.0_alpha_0.1.npy", history_resnet18)
```

### 4.3. Обучение без помощи учителя и сравнение результатов

Повторим процедуру предыдущего пункта, только без использования меток модели-учителя.

```python
os.makedirs("./experiments/hyperparams", exist_ok=True)
seed_everything(SEED)

resnet18 = resnet18_student2(N_CLASSES, DEVICE)
trainable = get_number_of_params(resnet18, trainable=True)
total = get_number_of_params(resnet18, trainable=False)
print(f"{trainable} (of {total}) trainable params")
train_dataset = ImageDataset(train_val_files, LABEL_ENCODER)
test_dataset = ImageDataset(test_files, LABEL_ENCODER)
opt = torch.optim.Adam(resnet18.parameters(), lr=0.0003)
criterion = nn.CrossEntropyLoss()
history_resnet18_no_teacher = train(train_dataset, test_dataset, resnet18, 50, 64, DEVICE, opt, criterion)
plot_training_history(history_resnet18_no_teacher)
np.save(f"./experiments/hyperparams/history_resnet18_no_teacher.npy", history_resnet18_no_teacher)
```

Построим графики accuracy моделей на тестовой выборке.

```python
history_resnet18_teacher = np.load("./experiments/hyperparams/history_resnet18_T_4.0_alpha_0.1.npy")
history_resnet18_no_teacher = np.load("./experiments/hyperparams/history_resnet18_no_teacher.npy")
plt.figure(figsize=(10, 6))
plt.plot(history_resnet18[:, 3], label = "T=4, alpha=0.1")
plt.plot(history_resnet18_no_teacher[:, 3], label = "no teacher")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.legend()
plt.grid()
```

Примерно к 30-50 эпохам обе модели выходят на плато. Судя по всему, KD дал небольшой прирост в метрике. Результат, конечно, не очень впечатляет. Можно сравнивать финальное качество, усредняя результаты последних 20-ти эпох.

```python
print("KD: {:.4f}, no teacher: {:.4f}".format(history_resnet18_teacher[-20:, 3].mean(), history_resnet18_no_teacher[-20:, 3].mean()))
```

    KD: 0.9302, no teacher: 0.9262

## 5. Эксперимент с разными температурами

Я провёл серию аналогичных экспериментов для $\alpha = 0.1$ и $T = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]$.

```python
temperatures = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]
means = []

plt.figure(figsize=(10, 6))
for t in temperatures:
    history = np.load(f"/content/experiments/hyperparams/history_resnet18_T_{t}_alpha_0.1.npy")
    plt.plot(history[:, 3], label = f"T={t}")
    means.append(history[-20:, 3].mean())

history_no_teacher = np.load("./experiments/hyperparams/history_resnet18_no_teacher.npy")
mean_no_teacher = history_no_teacher[-20:, 3].mean()
plt.plot(history_no_teacher[:, 3], label = "no teacher", color="black")

plt.xlabel("Epoch")
plt.ylabel("Test accuracy")
plt.title("Alpha = 0.1")
plt.ylim(0.91, 0.935)
plt.legend()
plt.grid()
plt.show()
```

```python
plt.figure(figsize=(10, 6))
plt.plot(temperatures, means, "bX")
plt.hlines(mean_no_teacher, 0, 9, "red", label="no teacher score")
plt.xlim(0, 9)
plt.xlabel("Temperature")
plt.ylabel("Mean test accuracy")
plt.title("Alpha = 0.1")
plt.legend()
plt.grid()
```

В этом эксперименте почти все значения температуры дали небольшой прирост в метрике. Оптимальная температура находится где-то в диапазоне от 3 до 6. Эффект от KD получился совсем небольшим. В [этом](https://github.com/peterliht/knowledge-distillation-pytorch#results-shallow-and-deep-distillation) репозитории проводились похожие эксперименты для задач Computer Vision (на датасете CIFAR-10), и улучшение качества тоже было незначительным.

## 6. Влияние размера обучающей выборки на качество дистилляции

Интересно посмотреть, как будут обучаться модели при уменьшении размера обучающей выборки. Я взял доли от исходного датасета $p = [0.01, 0.05, 0.1]$.

```python
ps = [0.01, 0.05, 0.1]

plt.figure(figsize=(10, 6))
for p in ps:
    history = np.load(f"./experiments/subset/history_resnet18_T_4.0_alpha_0.1_p_{p}.npy")
    plt.plot(history[:, 3], label = f"alpha = 0.1, T = 4.0, p = {p}")

    history_no_teacher = np.load(f"./experiments/subset/history_resnet18_no_teacher_p_{p}.npy")
    plt.plot(history_no_teacher[:, 3], label = f"no teacher, p = {p}")

plt.xlabel("Epoch")
plt.ylabel("Test accuracy")
plt.legend()
plt.grid()
plt.show()
```

KD лучше всего показал себя в самом экстремальном случае, когда обучение велось на $< 100$ картинках. При увеличении обучающей выборки модель получает возможность хорошо генерализоваться и эффект от KD вносит уже меньший вклад в финальное качество.

## 7. Промежуточные выводы

Результаты, которые я получил в экспериментах, показывают незначительное улучшение модели-студента при использовании KD. Этого всё ещё недостаточно для того, чтобы применять метод в реальных задачах. Тем не менее я бы хотел изучить работу knowledge distillation совместно с другими методами сжатия нейросетей. У меня есть гипотеза, что он может принести пользу совместно с [pruning](https://arxiv.org/pdf/1611.06440.pdf), где может быть задействован на стадии fine-tuning (возможно, получится снизить потерю качества после итерации pruning). Такие исследования уже [проводились](https://arxiv.org/pdf/1801.05787.pdf), и есть даже неплохие результаты.


# About

This is a demo of sampling-based minimum Bayes risk decoding for NMT [(Eikema and Aziz, 2020)](https://www.aclweb.org/anthology/2020.coling-main.398/) for educational purposes. For a scalable implementation check [mbr_nmt](https://github.com/Roxot/mbr-nmt) and [(Eikema and Aziz, 2021)](https://arxiv.org/abs/2108.04718).

# Setting Up

This part will take care of the following:

1. Clone the repository `demo-mbr-nmt`.
2. Install dependencies.
3. Download pre-trained models.

```python
!git clone https://github.com/probabll/demo-mbr-nmt.git
```

    fatal: destination path 'demo-mbr-nmt' already exists and is not an empty directory.

```python
cd demo-mbr-nmt
```

    /content/demo-mbr-nmt

```python
!git pull
```

    Already up to date.

```python
!pip install -r requirements.txt
```

    Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 1)) (1.9.0+cu102)
    Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 2)) (1.19.5)
    Requirement already satisfied: fairseq==0.10.2 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 3)) (0.10.2)
    Requirement already satisfied: sacrebleu==1.5.1 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 4)) (1.5.1)
    Requirement already satisfied: sacremoses==0.0.43 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 5)) (0.0.43)
    Requirement already satisfied: sentencepiece==0.1.95 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 6)) (0.1.95)
    Requirement already satisfied: regex in /usr/local/lib/python3.7/dist-packages (from fairseq==0.10.2->-r requirements.txt (line 3)) (2019.12.20)
    Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from fairseq==0.10.2->-r requirements.txt (line 3)) (4.62.0)
    Requirement already satisfied: dataclasses in /usr/local/lib/python3.7/dist-packages (from fairseq==0.10.2->-r requirements.txt (line 3)) (0.6)
    Requirement already satisfied: cffi in /usr/local/lib/python3.7/dist-packages (from fairseq==0.10.2->-r requirements.txt (line 3)) (1.14.6)
    Requirement already satisfied: cython in /usr/local/lib/python3.7/dist-packages (from fairseq==0.10.2->-r requirements.txt (line 3)) (0.29.24)
    Requirement already satisfied: hydra-core in /usr/local/lib/python3.7/dist-packages (from fairseq==0.10.2->-r requirements.txt (line 3)) (1.1.1)
    Requirement already satisfied: portalocker==2.0.0 in /usr/local/lib/python3.7/dist-packages (from sacrebleu==1.5.1->-r requirements.txt (line 4)) (2.0.0)
    Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from sacremoses==0.0.43->-r requirements.txt (line 5)) (1.15.0)
    Requirement already satisfied: click in /usr/local/lib/python3.7/dist-packages (from sacremoses==0.0.43->-r requirements.txt (line 5)) (7.1.2)
    Requirement already satisfied: joblib in /usr/local/lib/python3.7/dist-packages (from sacremoses==0.0.43->-r requirements.txt (line 5)) (1.0.1)
    Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch->-r requirements.txt (line 1)) (3.7.4.3)
    Requirement already satisfied: pycparser in /usr/local/lib/python3.7/dist-packages (from cffi->fairseq==0.10.2->-r requirements.txt (line 3)) (2.20)
    Requirement already satisfied: omegaconf==2.1.* in /usr/local/lib/python3.7/dist-packages (from hydra-core->fairseq==0.10.2->-r requirements.txt (line 3)) (2.1.1)
    Requirement already satisfied: antlr4-python3-runtime==4.8 in /usr/local/lib/python3.7/dist-packages (from hydra-core->fairseq==0.10.2->-r requirements.txt (line 3)) (4.8)
    Requirement already satisfied: importlib-resources in /usr/local/lib/python3.7/dist-packages (from hydra-core->fairseq==0.10.2->-r requirements.txt (line 3)) (5.2.2)
    Requirement already satisfied: PyYAML>=5.1.0 in /usr/local/lib/python3.7/dist-packages (from omegaconf==2.1.*->hydra-core->fairseq==0.10.2->-r requirements.txt (line 3)) (5.4.1)
    Requirement already satisfied: zipp>=3.1.0 in /usr/local/lib/python3.7/dist-packages (from importlib-resources->hydra-core->fairseq==0.10.2->-r requirements.txt (line 3)) (3.5.0)

```python
!bash de-en.sh
```

    Downloading de-en
    --2021-09-06 16:11:10--  https://surfdrive.surf.nl/files/index.php/s/8lnmJPX3iXndRpz/download
    Resolving surfdrive.surf.nl (surfdrive.surf.nl)... 145.100.27.67, 2001:610:108:203b:0:a11:da7a:5afe
    Connecting to surfdrive.surf.nl (surfdrive.surf.nl)|145.100.27.67|:443... connected.
    HTTP request sent, awaiting response... 200 OK
    Length: 1097465344 (1.0G) [application/x-tar]
    Saving to: ‘models/fairseq/de-en.tar’

    models/fairseq/de-e 100%[===================>]   1.02G  23.4MB/s    in 52s

    2021-09-06 16:12:02 (20.3 MB/s) - ‘models/fairseq/de-en.tar’ saved [1097465344/1097465344]

    Uncompressing
    ._de-en
    de-en/
    de-en/._prep-newstest2018.sh
    de-en/prep-newstest2018.sh
    de-en/._dict.de.txt
    de-en/dict.de.txt
    de-en/._.DS_Store
    de-en/.DS_Store
    de-en/._sentencepiece.bpe.model.de
    de-en/sentencepiece.bpe.model.de
    de-en/._sentencepiece.bpe.model.en
    de-en/sentencepiece.bpe.model.en
    de-en/._README.md
    de-en/README.md
    de-en/._sentencepiece.bpe.vocab.de
    de-en/sentencepiece.bpe.vocab.de
    de-en/._dict.en.txt
    de-en/dict.en.txt
    de-en/data/
    de-en/._sentencepiece.bpe.vocab.en
    de-en/sentencepiece.bpe.vocab.en
    de-en/._mle
    de-en/mle/
    de-en/._truecaser
    de-en/truecaser/
    de-en/truecaser/._truecase-model.en
    de-en/truecaser/truecase-model.en
    de-en/truecaser/._truecase-model.de
    de-en/truecaser/truecase-model.de
    de-en/mle/._averaged_model.pt
    de-en/mle/averaged_model.pt
    de-en/mle/._train.log
    de-en/mle/train.log
    de-en/data/._newstest2018.de-en.en
    de-en/data/newstest2018.de-en.en
    de-en/data/._news-commentary-v12.de-en.en
    de-en/data/news-commentary-v12.de-en.en
    de-en/data/._newstest2018.de-en.de
    de-en/data/newstest2018.de-en.de
    de-en/data/._news-commentary-v12.de-en.de
    de-en/data/news-commentary-v12.de-en.de
    Check: models/fairseq/de-en

# Neural machine translation

Given an input $x$ (i.e., a sentence in a source language), a trained NMT model predicts a conditional distribution $Y|X=x, \theta$ over all possible target-language translations of $x$.

The sample space $\mathcal Y$ of the model is the set of all sequences of the form $(y_1, \ldots, y_N)$ where $y_n$ belongs to a vocabulary of known target-language symbols, $N \ge 1$ is the sequence length, and $y_N$ is a special end-of-sequence (EOS) symbol.

An outcome $y \in \mathcal Y$ is assigned probability mass:

\begin{align}
P_{Y|X}(y|x,\theta) &= \prod_{n=1}^{N} \mathrm{Cat}(y_n|f(x, y_{<n}; \theta))
\end{align}

where $y_{<n}$ is the sequence of tokens before the $n$th token, and $f(\cdot; \theta)$ is a neural network architecture with parameters $\theta$.

In the previous section you downloaded a pre-trained NMT model, that is, a specific choice of architecture (a Transformer encoder-decoder model) and a collection of parameters $\theta$. This model was trained to maximise the model's likelihood given a dataset of observations (i.e., our training data).

You can load the trained model using `demo.load`:

```python
from demo import load
```

The next line will take a moment.

```python
model = load(src='de', tgt='en', np_seed=10, torch_seed=10)
```

Given an input, a trained model can assign probability to known translations, as well as, be used to draw translation candidates.

## Assigning probability

Let's first see how to compute the surprisal of a given translation (i.e, the negative logarithm of the outcome's probability):

```python
# Here is an example input in German and its human translation in English
example_x = "Es war noch nie leicht, ein rationales Gespräch über den Wert von Gold zu führen."
example_y = "It has never been easy to have a rational conversation about the value of gold."
```

```python
# Surprisal is the negative of the log-probability of the outcome
surp_obs = model.x2y.surprisal(
    example_x,
    example_y
)['surprisal']
print(surp_obs, (-surp_obs).exp())
```

    tensor([3.3471]) tensor([0.0352])

## Drawing samples

Now let's see how to draw translation candidates from the model distribution:

```python
for y in model.x2y.ancestral_sampling(example_x, num_samples=10)['output']:
    print(y)
```

    /usr/local/lib/python3.7/dist-packages/torch/_tensor.py:575: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.
    To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at  /pytorch/aten/src/ATen/native/BinaryOps.cpp:467.)
      return torch.floor_divide(self, other)

    It has never been easy to talk rationally about gold's value.
    It has never been easy to start a rational conversation about the value of gold.
    It has never been easy to make a rational talk about the value of gold.
    It has never been easy to begin a rational conversation on the value of gold.
    It has never been easy to address the value of gold in a rational manner.
    It has never been easy to handle a rational conversation on the value of gold.
    Discussing gold's value rationally has never been easy.
    It is never easy to call a rational discussion about gold's value.
    It is never easy to talk rationally about gold value.
    We always enjoy discussing the value of gold.

## Predicting a translation

Given an input sentence $x$, a trained model predicts an entire distribution $Y|X=x, \theta$ over the space $\mathcal Y$ of possible translations.

In order to pick a single translation candidate as our preferred translation, we need to introduce a criterion for decision making. Next, we recap two such criteria.

### Maximum-a-posterior decoding

MAP decoding tells us to pick the most probable translation.
\begin{align}
y^{\mathrm{mode}} &= \arg\max_{h \in \mathcal Y} ~ P_{Y|X}(h|x, \theta)
\end{align}

As the space $\mathcal Y$ is unbounded and the NMT model makes no Markov assumptions, this search is intractable. In practice, we approximate the decision rule searching through a beam of probable translations.

\begin{align}
y^{\mathrm{beam}} &= \arg\max_{h \in \mathrm{beam}(x)} ~  P_{Y|X}(h|x, \theta)
\end{align}

```python
beam = model.x2y.beam_search(example_x, beam_size=5)
for y, surp in zip(beam['output'], beam['surprisal']):
    print(f"{surp:.4f}\t{y}")
y_map = beam['output'][0]
```

    2.4119	It has never been easy to talk rationally about the value of gold.
    3.3470	It has never been easy to have a rational conversation about the value of gold.
    3.6803	It has never been easy to discuss the value of gold rationally.
    4.6698	It has never been easy to discuss the value of gold in a rational way.
    5.7409	It has never been easy to discuss the value of gold in a rational manner.

**Some things to note**

* Often, the beam search translation has much higher probability (lower suprisal) than the reference translation.
* Often, if you increase the `beam_size` parameter, beam search will find translations that are even more probable (even lower surprisal).

### Minimum Bayes risk decoding

MBR decoding tells us to pick the translation that has highest expected utility:

\begin{align}
y^{\text{mbr}} &= \arg\max_{h \in \mathcal Y} ~ \underbrace{\mathbb E[u(Y, h; x)|x, \theta]}_{=: \mu_u(h; x, \theta)} \\
&= \arg\max_{h \in \mathcal Y} ~ \sum_{y \in \mathcal Y} u(y, h; x) P_{Y|X}(y|x, \theta)
\end{align}

where a utility function $u(y, h; x)$ quantifies the benefit of choosing $h$ as the translation of $x$, when $y$ is the correct (or preferred) translation.

There are two sources of intractability in MBR decoding. First,  just like in MAP decoding, the search space (i.e., the sample space $\mathcal Y$) is unbounded. Second, for any given candidate translation $h$, the expected utility $\mu_u(h; x, \theta) = \mathbb E[u(Y, h; x)|x, \theta]$ is intractable to compute.

[Eikema and Aziz (2020)](https://www.aclweb.org/anthology/2020.coling-main.398/) propose to

1. approximate the hypothesis space by a tractable subset of hypotheses $\mathcal H(x)$ obtained by sampling from the model;
2. approximate the expected value using Monte Carlo (MC).

The decision rule becomes

\begin{align}
y^{\text{smbr}} &= \arg\max_{h \in \mathcal H(x)} ~ \underbrace{\frac{1}{S} \sum_{s=1}^S u(y^{(s)}, h; x)}_{=: \hat\mu_u(h; x, \theta)}
\end{align}

where $y^{(s)} \sim Y|X=x, \theta$ is a sample from the NMT model (samples can be drawn efficiently via ancestral sampling).

For much more on sampling-based MBR see also [Eikema and Aziz (2021)](https://arxiv.org/abs/2108.04718).

We have prepared a number of helpers for the demo.

```python
from demo import SampleFromNMT, CachedSampler, Utility, MBR
```

Let's start by picking a utility function. In this demo we will assign utility using ChrF, a metric for automatic assessment of translation quality that works by comparing a translation candidate (a hypothesis) to a reference translation.

```python
from sacrebleu import sentence_chrf
```

For modularity, we need to make chrf compatible with `demo.Utility`, so here is a thin wrapper for that:

```python
class ChrF(Utility):

    def __call__(self, src: str, hyp: str, ref: str) -> float:
        return sentence_chrf(hyp, [ref]).score  # note that chrf does not make use of the source sentence

chrf = ChrF()
```

We then need a space of translation candidates, which MBR will score and rank. It's rather natural to use the NMT model itself to obtain such a hypothesis space, for example, via ancestral sampling:

```python
# Duplicate translations do not change the result of MBR, so we discard duplicates
hyp_space = list(set(model.x2y.ancestral_sampling(example_x, num_samples=100)['output']))
len(hyp_space)
```

    78

MBR requires a sampler for MC estimation, here we construct one.

```python
nmt_sampler = SampleFromNMT(model.x2y, sample_size=100)
```

```python
nmt_sampler(example_x)
```

    ['It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to talk rationally about the value of gold.',
     'It has never been easy to have a rational conversation about the value of gold.',
     'It has never been easy to have a rational conversation about the value of gold.',
     'It has never been easy to have a rational conversation about the value of gold.',
     'It has never been easy to engage in a rational conversation about the value of gold.',
     'It has never been easy to discuss the value of gold rationally.',
     'It has never been easy to discuss the value of gold rationally.',
     'It has never been easy to discuss the value of gold rationally.',
     'It has never been easy to discuss the value of gold rationally.',
     'It has never been easy to discuss the value of gold rationally.',
     'It has never been easy to discuss the value of gold rationally.',
     'It has never been easy to enter into a rational conversation about the value of gold.',
     'It has never been easy to have a rational discussion about the value of gold.',
     'It has never been easy to have a rational conversation on the value of gold.',
     'It has never been easy to engage in rational conversation about the value of gold.',
     'It has never been easy to engage rationally in a conversation about the value of gold.',
     'It has never been easy to conduct a rational conversation on the value of gold.',
     'It has never been easy to conduct a rational discussion about the value of gold.',
     "It was never easy to have a rational conversation about gold's value.",
     'It has never been easy to engage in rational talk about the value of gold.',
     'It has never been easy to have a rational debate about the value of gold.',
     'It has never been easy to direct a rational conversation about the value of gold.',
     "It has never been easy to make a rational discussion about gold's value.",
     'It has never been easy to talk rationally about gold.',
     'It has never been easy to discuss the value of gold on a rational basis.',
     "It has never been easy to converse rationally about gold's value.",
     'It has never been easy to talk sensibly about the value of gold.',
     "It has never been easy to have a rational exchange on gold's value.",
     "It has never been easy to talk sensibly about gold's value.",
     'It has never been easy to have a rational exchange of minds about the value of gold.',
     "To have a rational conversation about gold's value has never been easy.",
     'It has never been easy, ever, to engage in a rational conversation about the value of gold.',
     'It has never been easy to craft a rational discussion about the value of gold.',
     'It has never been easy to proceed with a rational conversation about the value of gold.',
     'It has never been easy to come to a rational discussion about the value of gold.',
     '"It has never been easy to speak rationally about the value of gold.',
     'It has never been easy to engage in rational discussions of the value of gold.',
     'It has never always been easy to talk rationally about the value of gold.',
     'It has never been easy to enter a rational conversation over the value of gold.',
     'It has never been easy to talk rationally about the value of Gold.',
     'It has never been easy to carry out a rational talk about the value of gold.',
     'It has never been easy to have a rational conversation about the amount of gold.',
     'It has never been easy to make rational exchanges about the value of gold.',
     'It used to be easy to talk rationally about the value of gold.',
     'It has never been easy to drink a rational conversation about the value of gold.',
     "It has never been easy to engage rationally with talk about gold's value.",
     'Talking rationally about the value of gold was never easy.',
     'A rational exchange of views about the value of gold has never been easy.',
     'Living a rational conversation about the value of gold has never been easy.',
     '"It has never been easy to have a rational conversation about gold.',
     'It has never been easy to engage in rational management of the value of gold.',
     'It was never easy to have a rational chat about the value of gold.',
     'The value of gold has never been easy to talk rationally about.',
     'It has never been easy to hold a low-key conversation about the value of gold.',
     'Conducting rational talks about the value of gold has never been easy.',
     'It has never been easier to engage in rational talk of the value of gold.',
     'Measuring the value of gold has never been easy.',
     'It has never been easy to enter into a rational human conversation about the value of gold.',
     'There has never been a simple dialogue about the value of gold.',
     'It was never easy to talk rationally about the gold value.',
     'It was never easy to get any rational conversation about the value of gold.',
     'Making rational conversation about the value of gold has never been easy.',
     "Indeed, it's never made easy to talk rationally about the value of gold.",
     'It has never been easy to hold rational informed discussions about the value of gold.',
     'It has never been easy to clash in rational terms about the value of gold.',
     'It has never been easier to embark on rational talk about the value of gold.',
     'It has always been easier to converse rationally about the value of gold.',
     'Placing a rational conversation on the value of gold was never easy.',
     'To be reasonable about the value of gold has never been easy.',
     'It has never been easy to be coherent with the value of gold.',
     'It has never been simple to purposefully talk about the value of gold.',
     'It has never been easy to be rational in the eyes of a conversation about gold.',
     'It never was easy to engage in a rational conversation about gold-like values.',
     'Raising any rationality about the value of gold has never been easy.',
     'Conducting rational words about the value of gold has never been easy.',
     'At one point it has never been easy to discuss the real value of gold rationally.',
     'This was not easy to reach a rational conversation about the value of gold.',
     'It has never been easy to talk rationally on how well gold will be worth.',
     'It has never been easy to conduct rational debate about property & # x02BC; s value.',
     'No rational point being made about the value of gold has ever been easy.',
     'It has never been easy to achieve a rational taxation of gold.',
     'It has never been easy to tear rhetoric about the value of gold to crisis.',
     "It is never easy to bring anyone back to a rational point of view about gold's value.",
     'The way it is rationally spent on the value of gold has never been easy to discage.',
     'Raising comprehension about gold values has never been easy;',
     'There has never been a business touch on when anyone has talked rationally about gold values.',
     'A rational discussion about the gold value has never I had a mind come along.',
     'Anyone who had some information on gold buying a shelf like fluxes had the opportunity to view the buying price of gold in real companies with in mind.',
     'The European Union has been relying on a genuine individual for many years to reinforce our national political action and root causes.',
     '€? What budget was distributed for the profit from the gold granted?',
     'Proposes powdered financial benefits for beginning buyers. Small and medium-sized companies is going to purchase DLCGs, but the forced selling them till the end meant lending money to their customers integrated in the marketplace.']

Sampling on CPU may be considered an expensive operation, as in MBR we need to draw samples repeatedly for each candidate in the hypothesis space, we may consider using a single shared pool of samples for all candidates. `demo.CachedSampler` is a helper for that.

```python
cached_sampler = CachedSampler(nmt_sampler, sample_size=30)
```

```python
cached_sampler(example_x)
```

    ['It was never easy to discuss the value of gold rationally.',
     'Speaking of the value of gold has never been easy.',
     'It has never been easy to start through a rational conversation about the value of the gold.',
     'Bond value was never easy to breathe: It was never the easy moment to talk rationally about the value of gold.',
     'It has never been easy to make a rational conversation about currency values.',
     "It was never easy to talk rationally about gold's value.",
     'It has never been easy to discuss the value of gold in a rational manner.',
     'It has never been easy to discuss the value of gold more rationally.',
     'It has never been easy to discuss the value of gold more rationally.',
     'Any rational discussion on the value of gold has never been easy.',
     'Rational talk about a gold value has never been easy.',
     'Professor Guisenberg presided over an art test in 1952, designed by Thomas Altowski.',
     'It has never been easy to hold a rational conversation about currency values.',
     'It has never been easy to set a rational conversation about the value of gold.',
     'It has never been easy to discuss the value of gold more rationally.',
     'It has never been easy to enter into a rational conversation about the value of gold.',
     'It has never been easy to enter into a rational conversation about the value of gold.',
     'It has never been easy to discuss the value of gold in a rational way.',
     'It has never been easy to enter into a rational conversation about the value of gold.',
     'It has never been easy to keep a reasonable conversation about the value of gold.',
     'It has never been easy to discuss the value of gold rationally.',
     'A rational conversation about the value of gold has never been easy.',
     'Any rational discussion on the value of gold has never been easy.',
     'It has never been easy to enter into a rational conversation about the value of gold.',
     "It's never been simple to have a rational conversation about gold's value.",
     'It has never been easy to have a rational conversation on the value of gold.',
     'It has never been easy to discuss in a rational way the value of gold.',
     'It has never been easy to regard the value of gold in a rational fashion.',
     'It has never been easy to talk rationally about the value of gold.',
     "It has never been easy to engage in rational conversation about gold's cost.",
     'It has never been easy to make a rational conversation about the value of gold.',
     "A rational conversation about gold's value has never been easy.",
     'Today, the rapid have hardly been a group which honestly told, or the Spobiros are Québecians.',
     'It has never been easy to make a rational conversation about currency values.',
     'It has never been easy to make a rational conversation about currency values.',
     'It has never been easy to engage in a rational conversation about the value of gold.',
     'It has never been easy to input a rational conversation about price of gold.',
     "It's never easier to talk rationally about the value of gold.",
     'It has never been easy to talk rationally about gold.',
     'It has never been quick to engage in rational conversations about the value of gold.',
     'An informed discussion about the value of gold was never easy.',
     'It has never been easy to discuss the value of gold in a rational way.',
     'Professor Guisenberg presided over an art test in 1952, designed by Thomas Altowski.',
     'It has never been easy to conduct a rational conversation on the value of gold.',
     'It has never been easy to have a rational conversation about the value of gold.',
     'It has never been easy to start through a rational conversation about the value of the gold.',
     'Having a rational conversation about the value of gold has never been easy; this will probably mean the exchange of money.',
     'It has never been easy to regard the value of gold in a rational fashion.',
     'Discussions of a caliber of gold came too early.',
     'Discussions of a caliber of gold came too early.']

```python
mbr = MBR(chrf, cached_sampler)  # This version of MBR is what Eikema and Aziz (2021) call MBR N-by-S
```

```python
y_mbr = mbr.decode(example_x, hyp_space)
y_mbr
```

    'It has never been easy to make any rational conversation about the value of gold.'

```python
chrf(example_x, y_map, example_y), chrf(example_x, y_mbr, example_y)
```

    (0.6750826084658059, 0.8838063208456114)

```python
from tabulate import tabulate
```

```python
rows = []
for h, mu in sorted(zip(hyp_space, mbr.mus(example_x, hyp_space)), key=lambda pair: pair[1], reverse=True):
    quality = chrf(example_x, h, example_y)
    rows.append([mu, quality, h])
print(tabulate(rows, headers=['E[chrf(x, h, Y)]', 'chrf(x, h, ref)', 'candidate']))
```

      E[chrf(x, h, Y)]    chrf(x, h, ref)  candidate
    ------------------  -----------------  --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
              0.664658           0.886664  It has never been easy to conduct a rational conversation about the value of gold.
              0.659087           0.90062   It has never been easy to say a rational conversation about the value of gold.
              0.655118           1         It has never been easy to have a rational conversation about the value of gold.
              0.641811           0.90094   It has never been easy to do a rational conversation about the value of gold.
              0.637213           0.835729  It has never been easy to have a rational conversation about gold's value.
              0.633381           0.883806  It has never been easy to make any rational conversation about the value of gold.
              0.628271           0.911187  It has never been easy to hold a rational conversation about the value of gold.
              0.620876           0.817578  Conducting a rational conversation about the value of gold has never been easy.
              0.618684           0.721071  It has never been easy to engage in a rational discussion about the value of gold.
              0.617675           0.612032  It has never been easy to discuss the value of gold in a rational fashion.
              0.617666           0.670212  It has never been easy to enter into dialogue rationally about the value of gold.
              0.614589           0.60407   It has never been easy to discuss the value of gold in a rational way.
              0.613682           0.675083  It has never been easy to talk rationally about the value of gold.
              0.610339           0.886397  It has never been easy to thrill on a rational conversation about the value of gold.
              0.606124           0.884556  It has never been easy to have a rational conversation over the value of gold.
              0.605274           0.767775  It has never been easy to have a rational discussion about a meaning up to the value of gold.
              0.602842           0.678189  It has never been easy to speak rationally about the value of gold.
              0.596541           0.677154  It has never been easy to engage in rationality about the value of gold.
              0.592622           0.587346  It has never been easy to discuss the value of gold rationally.
              0.589643           0.714952  It has never been easy to conduct rational discussion about the value of gold.
              0.586124           0.709362  It has never been easy to direct rational discussion about the value of gold.
              0.583383           0.6227    It has never been easy to negotiate rationally on the value of gold.
              0.580751           0.695361  It has never been easy to enter into a rational debate about the value of gold.
              0.578215           0.618961  It has never been easy to engage in a rational dialogue on the value of gold.
              0.576175           0.678455  It has never been easy to engage in rational debate about the value of gold.
              0.573922           0.692405  It has never been easy to ask a rational word about the value of gold.
              0.573308           0.618745  It's never been easy to talk rationally about the value of gold.
              0.572303           0.62791   It has never been easy to be rational in the question of the value of gold.
              0.571277           0.504081  It has never been easy to talk rationally about gold's value.
              0.56936            0.650222  More generally, it has never been easy to speak rationally about the value of gold.
              0.56746            0.67801   It has never been easy to be rational when talking about the value of gold.
              0.564824           0.715683  It has never been easy to engage in a rational conversation about gold.
              0.564819           0.730348  It has never been easy to conduct a rational conversation about gold values.
              0.563938           0.677562  It has never been easy to engage in rational talk about the value of gold.
              0.561168           0.792581  It has never been so easy to guide rational conversations about the value of gold.
              0.560604           0.614493  But it has never been easy to enter into a rational debate on the value of gold.
              0.558383           0.748634  It has never been easy to hold a rational conversation about gold's value.
              0.553609           0.639809  It has never been easy to talk Rationally about the value of gold always.
              0.548913           0.610123  It has never been easy to engage rationally on the value of gold.
              0.54638            0.724392  It has never been easy to have a rational discussion of the value of gold.
              0.545544           0.504081  "It has never been easy to talk rationally about gold values.
              0.542669           0.488271  It has never been easy to talk rationally about gold.
              0.541175           0.813553  In order to sustain the quality, it has never been easy to have a rational conversation on the value of gold.
              0.538226           0.744892  It has always been easier to keep a rational conversation about the value of gold.
              0.537107           0.709403  Exchanging a rational conversation on the value of gold has never been easy.
              0.534938           0.562835  It's never been easy to discern the value of gold and be a rational proposer.
              0.527468           0.592258  It has never been easy to talk rationally about a value of gold.
              0.523324           0.540816  It has never been easy to engage in rational discussion about gold's value.
              0.506507           0.612627  It has never been easy to chat sensibly about the value of gold.
              0.499075           0.531628  It was never easy to talk rationally about the value of gold.
              0.496709           0.479668  It has never been easy to deal rationally with gold values.
              0.49528            0.50747   It has never been easy still to discuss systematically the value of gold.
              0.491318           0.764882  It has never been easy to have a rational conversation on gold's value.
              0.486457           0.752891  It has always been easy to have an informal conversation about the value of gold.
              0.486229           0.531629  Talking about the value of gold has never been easy.
              0.478857           0.775973  It was never easy to hold a rational conversation about the value of gold.
              0.476547           0.531508  It has never been easy to have rational talk on its value.
              0.474625           0.631752  It has never been easy to keep a rational conversation focused on what gold ought to mean.
              0.471237           0.424792  It has never been easy to discuss money rationally.
              0.467101           0.460624  It was never easy to discuss the value of gold in a rational way.
              0.461994           0.515633  It has never been easy to engage constructively with a given value of gold.
              0.460473           0.489473  It has never been easy to be rational about gold.
              0.439256           0.459534  It has never been easy to hold it rationally on incomes.
              0.437783           0.560683  Indeed, it has never been simple to hold rational conversations about gold.
              0.435147           0.464798  Discussing the value of gold with one another, universically speaking, has never been easy.
              0.432098           0.631468  It used never to be easy to have a rational mindset about the value of gold.
              0.431282           0.489732  "Gold has never been easy to make a rational remark on what value it has to offer.
              0.409525           0.435477  It's never been easy to put amounts to gold round about money in a rational way.
              0.40589            0.451296  Discussing the value of gold has never been easy.
              0.383871           0.403844  It's never been easy to start describing gold in a rational way.
              0.383609           0.448246  One could never have seen a rational discussion about a value of gold facilitating this.
              0.373323           0.358997  It is never easy or simply any rational discussion about gold's value.
              0.35126            0.341711  To move independently, a bilateral transaction between the two was better. There was no time to discuss the value of the gold rationally.
              0.336515           0.373002  It was never easy to be rational about gold with value.
              0.281089           0.329801  Using energy for rational amendment of investors' price _ rumours and as services are reducing the effort linked with conversion, prices and exports have never been Frankfurt.
              0.256596           0.261081  It's never a problem management talk about gold's value, as they reported a major theory - whether women like gold prizes or half a dollar awards - is uncomplicated.
              0.236358           0.258179  Pragmatism and equivalence are both the objective and the test of mind. An incapable description is however: Rational meaning is dismissed as not the substance of a proper peg.
              0.120835           0.100592  Getting to know gold is best possible.

```python

```

### Analysis

```python
import numpy as np
import matplotlib.pyplot as plt
```

Here we analyse whether a candidate's expected utility $\mu_u(h;x, \theta)$ correlates with utility against the ground-truth. As $\mu_u(h;x, \theta)$ is intractable to  compute, we rely on stochastic estimates $\hat\mu_u(h;x, \theta)$.

```python
from tqdm.auto import tqdm
```

We repeat the next experiment a number of times to account for variance, which may take some time.

```python
estimates = []
# compute utility against reference
for h in hyp_space:
    cached_sampler.reset()  #
    u = chrf(example_x, h, example_y)  # use human reference
    estimates.append([u])
repetitions = 10  # let's repeat MC estimation a number of times to get a sense of variance
for _ in tqdm(range(repetitions)):
    cached_sampler.reset()  # reset cache to obtain a better view of variance
    for n, h in enumerate(hyp_space):
        estimates[n].append(mbr.mu(example_x, h))
estimates = np.array(estimates)
```

      0%|          | 0/10 [00:00<?, ?it/s]

```python
estimates.shape
```

    (78, 11)

```python
_ = plt.errorbar(
    estimates[:,0], estimates[:,1:].mean(1), yerr=estimates[:,1:].std(1) * 2,
    marker='.', linestyle='', color='gray', label=r'mean and var of $\hat\mu_u(h;\theta)$')
_ = plt.title("Estimates of expected utility vs ground-truth utility")
_ = plt.xlabel(r'$u(y_{ref}, h; x)$')
_ = plt.ylabel(r'$\hat \mu_u(h; x, \theta)$')
_ = plt.legend()
```

```python

```