Distributed Associative Memory Network with Association Reinforcing Loss
===

Requirements
---
  * CUDA 9.0
  * CUDNN 7
  * python 2.7
  * tensorflow 1.12
  * dm-sonnet 1.34
  
```setup
pip install -r requirements.txt
```
  
<br>

Training
---

After installing every required libraries, each task can be traind by below command. (DAM-2 version) <br>


* ####Representation Recall Task
```shell
python train.py \
  --mode=RepresentationRecall \
  --babi=False \
  --dam=True \
  --num_memory_blocks=2 \
  --batch_size=16 \
  --learning_rate=1e-4 \
  --hidden_size=128 \
  --memory_address_size=32 \
  --memory_length_size=60 \
  --num_read_heads=1 \
  --N=4 \
  --bit_w=64 \
  --num_bit=8
  --min_length=8 \
  --max_length=16 \
  --training_iteration=20000 \
  --name=path/to/checkpoint_dir
```


* ####Copy Task
```shell
python train.py \
  --mode=Copy \
  --babi=False \
  --dam=True \
  --num_memory_blocks=2 \
  --p_re=0.1 \
  --batch_size=16 \
  --learning_rate=1e-4 \
  --hidden_size=128 \
  --memory_address_size=64 \
  --memory_length_size=36 \
  --num_read_heads=1 \
  --bit_w=8 \
  --min_length=8 \
  --max_length=32 \
  --training_iteration=10000 \
  --name=path/to/checkpoint_dir
```


* ####Associative Recall Task
```shell
python train.py \
  --mode=AssociativeRecall \
  --babi=False \
  --dam=True \
  --num_memory_blocks=2 \
  --p_re=0.1 \
  --batch_size=16 \
  --learning_rate=1e-4 \
  --hidden_size=128 \
  --memory_address_size=32 \
  --memory_length_size=36 \
  --num_read_heads=1 \
  --bit_w=8 \
  --min_length=2 \
  --max_length=8 \
  --item_bit=3 \
  --training_iteration=10000 \
  --name=path/to/checkpoint_dir
```


* ####bAbI Task
  * Training
  ```shell
    python train.py \
      --babi=True \
      --dam=True \
      --num_memory_blocks=2 \
      --p_re=0.1 \
      --batch_size=32 \
      --learning_rate=3e-5 \
      --hidden_size=256 \
      --memory_address_size=128 \
      --memory_length_size=48 \
      --num_read_heads=4 \
      --epoch=50 \
      --name=path/to/checkpoint_dir
  ```
  
  * Fine-tuning
  ```shell
    python train.py \
      --babi=True \
      --dam=True \
      --num_memory_blocks=2 \
      --p_re=0.1 \
      --batch_size=32 \
      --learning_rate=1e-5 \
      --hidden_size=256 \
      --memory_address_size=128 \
      --memory_length_size=48 \
      --num_read_heads=4 \
      --epoch=5 \
      --name=path/to/checkpoint_dir
  ```

<br> 

Evaluation
---

* ####bAbI Task

```shell
python eval.py \
  --dam=True \
  --num_memory_blocks=2 \
  --hidden_size=256 \
  --memory_address_size=128 \
  --memory_length_size=48 \
  --num_read_heads=4 \
  --name=path/to/checkpoint_dir \
  --num=<the number of training iterations>
```
  
<br>


```shell
python eval.py \
  --dam=True \
  --num_memory_blocks=2 \
  --hidden_size=256 \
  --memory_address_size=128 \
  --memory_length_size=48 \
  --num_read_heads=4 \
  --name=model \
  --num=107360
```


Acknowledement
---
Code based on [DeepMind's DNC](https://github.com/deepmind/dnc)