<jupyter_start><jupyter_text>A Design Space Study for LISTA and Beyond---Copyright 2021 Google Research, Licensed under the Apache License, Version 2.0  0. Install package and download checkpoints/dictionary.<jupyter_code>%%capture
#@title
import os
import sys

# Download source code.
if "lista_design_space" not in os.getcwd():
  !git clone --depth 1 https://github.com/google-research/google-research
  os.chdir('google-research/lista_design_space')
  sys.path.append('.')
  !pip install -r requirements.txt
else:
  !git pull

os.mkdir('data')
!wget https://storage.googleapis.com/gresearch/lista_design_space/A_250_500.npy -P data
os.mkdir('checkpoints')
!wget https://storage.googleapis.com/gresearch/lista_design_space/lista.h5 -P checkpoints
!wget https://storage.googleapis.com/gresearch/lista_design_space/lfista.h5 -P checkpoints
!wget https://storage.googleapis.com/gresearch/lista_design_space/dense.h5 -P checkpoints
!wget https://storage.googleapis.com/gresearch/lista_design_space/searched.h5 -P checkpoints<jupyter_output><empty_output><jupyter_text>1. Generate test data.<jupyter_code>!python data_generator.py<jupyter_output>2021-03-31 23:15:45.169813: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0<jupyter_text>2. Define loss function and metrics.<jupyter_code>import os
import sys

import numpy as np
import tensorflow as tf
import lista_models


class MSE(tf.keras.losses.Loss):
  """MSE loss for LISTA."""

  def __init__(self, dict_shape, **kwargs):
    super(MSE, self).__init__(**kwargs)
    self.dict_shape = dict_shape

  def call(self, y_true, y_pred):
    loss = tf.nn.l2_loss(y_pred[:, -self.dict_shape[1]:] - y_true)
    return loss


class NMSE(tf.keras.metrics.Mean):
  """NMSE metrics for LISTA."""

  def __init__(self, name, dict_shape, dtype=None):
    super(NMSE, self).__init__(name=name, dtype=dtype)
    self.dict_shape = dict_shape

  def update_state(self, y_true, y_pred, sample_weight=None):
    mse = tf.reduce_mean(
        tf.square(y_true - y_pred[:, -self.dict_shape[1]:]), axis=-1)
    nmse_denom = tf.reduce_mean(tf.square(y_true), axis=-1)
    nmse = mse / nmse_denom
    nmse_val = 10.0 * tf.math.log(nmse) / tf.math.log(10.0)
    return super(NMSE, self).update_state(nmse_val, sample_weight=sample_weight)<jupyter_output><empty_output><jupyter_text>3. Define input dataset.<jupyter_code>def input_fn(mode, batch_size, dict_shape):
  """Input dataset for LISTA."""
  filename = 'sc_{}_data.npy'.format(mode)
  data = np.load(
      os.path.join('./data', filename),
      allow_pickle=True).astype(np.float32)
  dataset = tf.data.Dataset.from_tensor_slices(data)

  def _parser(value):
    return value[:dict_shape[0]], value[dict_shape[0]:]

  dataset = dataset.map(_parser)
  dataset = dataset.batch(batch_size)
  dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

  return dataset<jupyter_output><empty_output><jupyter_text>4. Define test code.<jupyter_code>_NUM_LAYERS = 16
_LAMBDA = 0.4
_NUM_TEST_IMAGES = 10240
_TEST_BATCH_SIZE = 1024

def test(model_name):
  if model_name == 'lista':
    arch_str = '0_0_0_0_0_0_0_0_0_0_0_0_0_0_0_0'
  elif model_name == 'lfista':
    arch_str = '0_0_1_1_1_1_1_1_1_1_1_1_1_1_1_1'
  elif model_name == 'dense':
    arch_str = '0_0_1_3_7_15_31_63_127_255_511_1023_2047_4095_8191_16383'
  elif model_name == 'searched':
    arch_str = '0_0_0_2_1_3_14_42_48_249_284_4_152_328_4888_12193'
  else:
    raise ValueError('Model {} is not supported.'.format(model_name))

  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
  tf.config.experimental_connect_to_cluster(tpu)
  tf.tpu.experimental.initialize_tpu_system(tpu)
  strategy = tf.distribute.TPUStrategy(tpu)

  with strategy.scope():
    dictionary = np.load(
        tf.io.gfile.GFile('./data/A_250_500.npy', 'rb'),
        allow_pickle=True).astype(np.float32)
    test_dataset = input_fn('test', _TEST_BATCH_SIZE, dictionary.shape)
    model = lista_models.Lista(dictionary, _LAMBDA, arch_str, _NUM_LAYERS)
    model.build([None, dictionary.shape[0]])
    model.load_weights('./checkpoints/{}.h5'.format(model_name))

    model.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss=MSE(dictionary.shape),
        metrics=[NMSE('nmse', dictionary.shape)])
    model.evaluate(
        x=test_dataset,
        steps=_NUM_TEST_IMAGES // _TEST_BATCH_SIZE,
        verbose=2)<jupyter_output><empty_output><jupyter_text>5. Run test.<jupyter_code>test('lista')
test('lfista')
test('dense')
test('searched')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13093709884467267, 'fraction_numerical': 0.04150620453572957, 'mean_word_length': 4.993589743589744, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 24, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '543405', 'license': 'Apache-2.0', 'n_tokens_mistral': 1865, 'n_tokens_neox': 1757, 'n_words': 284, 'path': 'lista_design_space/lista_test.ipynb', 'repo_name': 'wy-go/google-research', 'size': '4674'}