<filename>core/lib/test_trainer.py<gh_stars>1-10
# Copyright (C) 2021 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import jax
import pytest
import unittest

from config import default as config_lib
from core.data import codenet_paths
from core.data import data_io
from core.data import info as info_lib
from core.lib import trainer

import tensorflow_datasets as tfds


def validate_single_step(config):
  info = info_lib.get_test_info()
  dataset_path = codenet_paths.TEST_DATASET_PATH
  split = 'train'
  steps = 1
  trainer.Trainer(config=config, info=info).run_train(
      dataset_path=dataset_path, split=split, steps=steps)


class TrainerTest(unittest.TestCase):

  @pytest.mark.slow
  def test_ipagnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.raise_in_ipagnn = True
    config.max_tokens = 512
    validate_single_step(config)

  # Disabled for timeout.
  # @pytest.mark.slow
  def test_ggnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'GGNN'
    config.batch_size = 16
    config.max_tokens = 512
    validate_single_step(config)

  def test_finetune_from_lstm(self):
    config = config_lib.get_test_config()
    config.model_class = 'LSTM'
    config.batch_size = 16
    config.max_tokens = 512
    config.experiment_id = 'tests-lstm'
    validate_single_step(config)

    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.max_tokens = 512
    config.finetune = 'LSTM'
    config.restore_checkpoint_dir = 'out/experiments/tests-lstm/checkpoints'
    validate_single_step(config)

  @pytest.mark.slow
  def test_film_ipagnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.use_film = True
    config.raise_in_ipagnn = False
    config.eval_freq = 1
    config.max_tokens = 512
    validate_single_step(config)

  @pytest.mark.slow
  def test_film_exception_ipagnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.use_film = True
    config.raise_in_ipagnn = True
    config.modulate_mode = 'concat'
    config.eval_freq = 1
    config.max_tokens = 512
    validate_single_step(config)

  @pytest.mark.slow
  def test_compressive_ipagnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.use_compressive_ipagnn = True
    config.eval_freq = 1
    config.compressive_max_skip = 3
    config.max_tokens = 512
    validate_single_step(config)


if __name__ == '__main__':
  unittest.main()

<filename>core/lib/test_trainer.py<gh_stars>1-10
# Copyright (C) 2021 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import jax
import pytest
import unittest

from config import default as config_lib
from core.data import codenet_paths
from core.data import data_io
from core.data import info as info_lib
from core.lib import trainer

import tensorflow_datasets as tfds


def validate_single_step(config):
  info = info_lib.get_test_info()
  dataset_path = codenet_paths.TEST_DATASET_PATH
  split = 'train'
  steps = 1
  trainer.Trainer(config=config, info=info).run_train(
      dataset_path=dataset_path, split=split, steps=steps)


class TrainerTest(unittest.TestCase):

  @pytest.mark.slow
  def test_ipagnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.raise_in_ipagnn = True
    config.max_tokens = 512
    validate_single_step(config)

  # Disabled for timeout.
  # @pytest.mark.slow
  def test_ggnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'GGNN'
    config.batch_size = 16
    config.max_tokens = 512
    validate_single_step(config)

  def test_finetune_from_lstm(self):
    config = config_lib.get_test_config()
    config.model_class = 'LSTM'
    config.batch_size = 16
    config.max_tokens = 512
    config.experiment_id = 'tests-lstm'
    validate_single_step(config)

    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.max_tokens = 512
    config.finetune = 'LSTM'
    config.restore_checkpoint_dir = 'out/experiments/tests-lstm/checkpoints'
    validate_single_step(config)

  @pytest.mark.slow
  def test_film_ipagnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.use_film = True
    config.raise_in_ipagnn = False
    config.eval_freq = 1
    config.max_tokens = 512
    validate_single_step(config)

  @pytest.mark.slow
  def test_film_exception_ipagnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.use_film = True
    config.raise_in_ipagnn = True
    config.modulate_mode = 'concat'
    config.eval_freq = 1
    config.max_tokens = 512
    validate_single_step(config)

  @pytest.mark.slow
  def test_compressive_ipagnn(self):
    config = config_lib.get_test_config()
    config.model_class = 'IPAGNN'
    config.batch_size = 16
    config.use_compressive_ipagnn = True
    config.eval_freq = 1
    config.compressive_max_skip = 3
    config.max_tokens = 512
    validate_single_step(config)


if __name__ == '__main__':
  unittest.main()

# Copyright (C) 2021 Google Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Disabled for timeout. # @pytest.mark.slow