<reponame>Slowika/GameBias-EmeCom2020<filename>lib/EGG_research/egg/zoo/summation/train.py
# Copyright (c) Facebook, Inc. and its affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from __future__ import print_function
import argparse
import torch.utils.data
import torch.nn.functional as F
import egg.core as core
from egg.zoo.summation.features import SequenceLoader
from egg.zoo.summation.archs import Encoder, Receiver


def get_params():
    parser = argparse.ArgumentParser()
    parser.add_argument('--batches_per_epoch', type=int, default=1000,
                        help='Number of batches per epoch (default: 1000)')

    parser.add_argument('--sender_hidden', type=int, default=10,
                        help='Size of the hidden layer of Sender (default: 10)')
    parser.add_argument('--receiver_hidden', type=int, default=10,
                        help='Size of the hidden layer of Receiver (default: 10)')

    parser.add_argument('--sender_embedding', type=int, default=10,
                        help='Dimensionality of the embedding hidden layer for Sender (default: 10)')
    parser.add_argument('--receiver_embedding', type=int, default=10,
                        help='Dimensionality of the embedding hidden layer for Receiver (default: 10)')

    parser.add_argument('--sender_cell', type=str, default='rnn',
                        help='Type of the cell used for Sender {rnn, gru, lstm} (default: rnn)')
    parser.add_argument('--receiver_cell', type=str, default='rnn',
                        help='Type of the cell used for Receiver {rnn, gru, lstm} (default: rnn)')

    parser.add_argument('--temperature', type=float, default=1.0,
                        help="GS temperature for the sender (default: 1.0)")
    parser.add_argument('--max_n', type=int, default=10,
                        help="Max n in a^nb^n(default: 10)")
    args = core.init(parser)

    return args


def loss(_sender_input, _message, _receiver_input, receiver_output, labels):
    acc = (receiver_output.argmax(dim=1) == labels).detach().float()
    loss = F.cross_entropy(receiver_output, labels)
    return loss, {'acc': acc}


if __name__ == "__main__":
    opts = get_params()

    device = torch.device("cuda" if opts.cuda else "cpu")

    train_loader = SequenceLoader(max_n=opts.max_n, batch_size=opts.batch_size,
                                  batches_per_epoch=opts.batches_per_epoch)
    test_loader = SequenceLoader(max_n=opts.max_n, batch_size=opts.batch_size,
                                 batches_per_epoch=opts.batches_per_epoch, seed=7)

    encoder = Encoder(n_hidden=opts.sender_hidden, embed_dim=opts.sender_embedding,
                     cell=opts.sender_cell, vocab_size=3)  # only 3 symbols in the incoming data
    sender = core.RnnSenderGS(encoder, opts.vocab_size, opts.sender_embedding, opts.sender_hidden,
                              cell=opts.sender_cell, max_len=opts.max_len, temperature=opts.temperature)

    receiver = Receiver(opts.receiver_hidden)
    receiver = core.RnnReceiverGS(receiver, opts.vocab_size, opts.receiver_embedding,
                                  opts.receiver_hidden, cell=opts.receiver_cell)

    game = core.SenderReceiverRnnGS(sender, receiver, loss)

    optimizer = core.build_optimizer(game.parameters())

    trainer = core.Trainer(game=game, optimizer=optimizer, train_data=train_loader,
                           validation_data=test_loader)
    trainer.train(n_epochs=opts.n_epochs)

    sender_inputs, messages, _, receiver_outputs, labels = \
        core.dump_sender_receiver(game, test_loader, gs=True, device=device, variable_length=True)

    for (seq, l), message, output, label in zip(sender_inputs, messages, receiver_outputs, labels):
        print(f'{seq[:l]} -> {message} -> {output.argmax()} (label = {label})')

    core.close()


# Copyright (c) Facebook, Inc. and its affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from __future__ import print_function
import argparse
import torch.utils.data
import torch.nn.functional as F
import egg.core as core
from egg.zoo.summation.features import SequenceLoader
from egg.zoo.summation.archs import Encoder, Receiver


def get_params():
    parser = argparse.ArgumentParser()
    parser.add_argument('--batches_per_epoch', type=int, default=1000,
                        help='Number of batches per epoch (default: 1000)')

    parser.add_argument('--sender_hidden', type=int, default=10,
                        help='Size of the hidden layer of Sender (default: 10)')
    parser.add_argument('--receiver_hidden', type=int, default=10,
                        help='Size of the hidden layer of Receiver (default: 10)')

    parser.add_argument('--sender_embedding', type=int, default=10,
                        help='Dimensionality of the embedding hidden layer for Sender (default: 10)')
    parser.add_argument('--receiver_embedding', type=int, default=10,
                        help='Dimensionality of the embedding hidden layer for Receiver (default: 10)')

    parser.add_argument('--sender_cell', type=str, default='rnn',
                        help='Type of the cell used for Sender {rnn, gru, lstm} (default: rnn)')
    parser.add_argument('--receiver_cell', type=str, default='rnn',
                        help='Type of the cell used for Receiver {rnn, gru, lstm} (default: rnn)')

    parser.add_argument('--temperature', type=float, default=1.0,
                        help="GS temperature for the sender (default: 1.0)")
    parser.add_argument('--max_n', type=int, default=10,
                        help="Max n in a^nb^n(default: 10)")
    args = core.init(parser)

    return args


def loss(_sender_input, _message, _receiver_input, receiver_output, labels):
    acc = (receiver_output.argmax(dim=1) == labels).detach().float()
    loss = F.cross_entropy(receiver_output, labels)
    return loss, {'acc': acc}


if __name__ == "__main__":
    opts = get_params()

    device = torch.device("cuda" if opts.cuda else "cpu")

    train_loader = SequenceLoader(max_n=opts.max_n, batch_size=opts.batch_size,
                                  batches_per_epoch=opts.batches_per_epoch)
    test_loader = SequenceLoader(max_n=opts.max_n, batch_size=opts.batch_size,
                                 batches_per_epoch=opts.batches_per_epoch, seed=7)

    encoder = Encoder(n_hidden=opts.sender_hidden, embed_dim=opts.sender_embedding,
                     cell=opts.sender_cell, vocab_size=3)  # only 3 symbols in the incoming data
    sender = core.RnnSenderGS(encoder, opts.vocab_size, opts.sender_embedding, opts.sender_hidden,
                              cell=opts.sender_cell, max_len=opts.max_len, temperature=opts.temperature)

    receiver = Receiver(opts.receiver_hidden)
    receiver = core.RnnReceiverGS(receiver, opts.vocab_size, opts.receiver_embedding,
                                  opts.receiver_hidden, cell=opts.receiver_cell)

    game = core.SenderReceiverRnnGS(sender, receiver, loss)

    optimizer = core.build_optimizer(game.parameters())

    trainer = core.Trainer(game=game, optimizer=optimizer, train_data=train_loader,
                           validation_data=test_loader)
    trainer.train(n_epochs=opts.n_epochs)

    sender_inputs, messages, _, receiver_outputs, labels = \
        core.dump_sender_receiver(game, test_loader, gs=True, device=device, variable_length=True)

    for (seq, l), message, output, label in zip(sender_inputs, messages, receiver_outputs, labels):
        print(f'{seq[:l]} -> {message} -> {output.argmax()} (label = {label})')

    core.close()
# Copyright (c) Facebook, Inc. and its affiliates. # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. # only 3 symbols in the incoming data