test_CSamplerWatcher.py

import tensorflow as tf
from Utils.utils import CFakeObject
from NN.restorators.samplers import sampler_from_config
from NN.restorators.samplers.CSamplerWatcher import CSamplerWatcher
from NN.utils import masked

def _fake_sampler(stochasticity=1.0, timesteps=10):
  interpolant = sampler_from_config({
    'name': 'DDIM',
    'stochasticity': stochasticity,
    'noise stddev': 'zero',
    'schedule': {
      'name': 'discrete',
      'beta schedule': 'linear',
      'timesteps': timesteps,
    },
    'steps skip type': { 'name': 'uniform', 'K': 1 },
  })
  
  shape = (32, 3)
  fakeNoise = tf.random.normal(shape)
  def fakeModel(V, T, **kwargs):
    return fakeNoise + tf.cast(T, tf.float32) * V

  x = tf.random.normal(shape)
  return CFakeObject(x=x, model=fakeModel, interpolant=interpolant)

def _fake_AR(threshold, timesteps=10, scale=1.0):
  interpolant = sampler_from_config({
    "name": "autoregressive",
    "noise provider": "zero",
    "threshold": threshold,
    "steps": {
      "start": 1.0,
      "end": 0.001,
      "steps": timesteps,
      "decay": 0.9
    },
    "interpolant": { "name": "direction" }
  })
 
  shape = (32, 3)
  fakeNoise = tf.random.normal(shape)
  def fakeModel(V, T, mask, **kwargs):
    s = fakeNoise
    if mask is not None:
      s = masked(fakeNoise, mask)
      T = masked(T, mask)
      V = masked(V, mask)
    
    return s + tf.cast(T, tf.float32) * V * scale

  x = tf.random.normal(shape)
  return CFakeObject(x=x, model=fakeModel, interpolant=interpolant)

def test_notAffectResults():
  fake = _fake_sampler()
  X_withoutWatcher = fake.interpolant.sample(value=fake.x, model=fake.model)

  watcher = CSamplerWatcher(steps=10, tracked=dict())
  X_withWatcher = fake.interpolant.sample(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())

  tf.debugging.assert_equal(X_withoutWatcher, X_withWatcher)
  return

def _commonChecks(collectedSteps, x, results):
  tf.debugging.assert_equal(tf.shape(collectedSteps)[1:], tf.shape(x))
  tf.debugging.assert_equal(collectedSteps[0], x, 'First step must be equal to initial value')
  tf.debugging.assert_equal(collectedSteps[-1], results, 'Last step must be equal to results')
  return

def test_collectsSteps():
  fake = _fake_sampler()
  watcher = CSamplerWatcher(
    steps=10,
    tracked=dict(
      value=(32, 3)
    )
  )
  oldX = watcher.tracked('value').numpy()
  results = fake.interpolant.sample(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())

  collectedSteps = watcher.tracked('value')
  tf.debugging.assert_greater(
    tf.reduce_sum(tf.abs(collectedSteps - oldX)), 0.0,
    'Tracked steps must be different from initial value'
  )
  _commonChecks(collectedSteps, fake.x, results)
  tf.debugging.assert_equal(tf.shape(collectedSteps)[0], 11, 'Must collect 11 values')
  tf.debugging.assert_equal(watcher.iteration, 10, 'Must collect 10 steps')
  return

def test_collectsOnlyIndicedValues():
  fake = _fake_sampler()
  indices = [0, 2, 6]
  watcher = CSamplerWatcher(
    steps=10,
    tracked=dict(
      value=(3,)
    ),
    indices=indices
  )
  results = fake.interpolant.sample(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())

  collectedSteps = watcher.tracked('value')
  _commonChecks(
    collectedSteps, 
    x=tf.gather(fake.x, indices, axis=0),
    results=tf.gather(results, indices, axis=0)
  )
  tf.debugging.assert_equal(tf.shape(collectedSteps)[0], 11, 'Must collect 11 values')
  return

def test_resetIteration():
  watcher = CSamplerWatcher(steps=10, tracked=dict())
  fake = _fake_sampler(timesteps=10)
  _ = fake.interpolant.sample(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())
  tf.debugging.assert_equal(watcher.iteration, 10, 'Must collect 10 steps')
  
  fake = _fake_sampler(timesteps=5)
  _ = fake.interpolant.sample(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())
  tf.debugging.assert_equal(watcher.iteration, 5, 'Must collect 5 steps')
  return

def _checkTracked(value, N):
  assert value is not None, 'Must be tracked'
  tf.debugging.assert_equal(tf.shape(value), (N, 32, 3), 'Unexpected shape')
  tf.debugging.assert_greater(tf.reduce_sum(tf.abs(value[0] - value[1])), 0.0, 'Must be different')
  return

def test_trackSolution():
  fake = _fake_sampler()
  watcher = CSamplerWatcher(
    steps=10,
    tracked=dict(x0=(32, 3), x1=(32, 3), value=(32, 3))
  )
  _ = fake.interpolant.sample(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())
  _checkTracked(watcher.tracked('x0'), N=10)
  _checkTracked(watcher.tracked('x1'), N=10)
  _checkTracked(watcher.tracked('value'), N=11)
  return

def test_trackSolutionWithMask():
  fake = _fake_AR(threshold=0.1)
  watcher = CSamplerWatcher(
    steps=10,
    tracked=dict(x0=(32, 3), x1=(32, 3))
  )
  _ = fake.interpolant.sample(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())
  _checkTracked(watcher.tracked('x0'), N=10)
  _checkTracked(watcher.tracked('x1'), N=10)
  return

def test_trackSolutionWithMask_value():
  fake = _fake_AR(threshold=0.5)
  watcher = CSamplerWatcher(
    steps=10,
    tracked=dict(value=(32, 3), x0=(32, 3), x1=(32, 3))
  )
  _ = fake.interpolant.sample(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())
  _checkTracked(watcher.tracked('value'), N=11)
  _checkTracked(watcher.tracked('x0'), N=10)
  _checkTracked(watcher.tracked('x1'), N=10)
  return

# test multiple calls with index
def test_multipleCallsWithIndex():
  fake = _fake_sampler()
  watcher = CSamplerWatcher(
    steps=10,
    tracked=dict(value=(32*3, 3))
  )
  arg = dict(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())
  A = fake.interpolant.sample(**arg, index=0)
  B = fake.interpolant.sample(**arg, index=32)
  C = fake.interpolant.sample(**arg, index=64)

  tf.debugging.assert_equal(A, B)
  tf.debugging.assert_equal(A, C)

  collectedSteps = watcher.tracked('value')
  tf.debugging.assert_equal(tf.shape(collectedSteps)[1], 96, 'Must collect 96 values')
  tf.debugging.assert_equal(watcher.iteration, 10, 'Must collect 10 steps')
  # values must be same across (0..32), (32..64), (64..96)
  tf.debugging.assert_equal(collectedSteps[:, :32], collectedSteps[:, 32:64])
  tf.debugging.assert_equal(collectedSteps[:, 32:64], collectedSteps[:, 64:])
  return

# test multiple calls with index and mask
def test_multipleCallsWithIndexAndMask():
  fake = _fake_AR(threshold=0.1)
  watcher = CSamplerWatcher(
    steps=10,
    tracked=dict(value=(3,)),
    indices=[0, 32, 64, 65]
  )
  arg = dict(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())
  A = fake.interpolant.sample(**arg, index=0)
  B = fake.interpolant.sample(**arg, index=32)
  C = fake.interpolant.sample(**arg, index=64)
  tf.assert_equal(A, B)
  tf.assert_equal(A, C)

  collectedSteps = watcher.tracked('value')[:watcher.iteration]
  tf.assert_equal(tf.shape(collectedSteps)[1:], (4, 3), 'Must be (4, 3)')
  tf.debugging.assert_equal(collectedSteps[:, 0:1], collectedSteps[:, 1:2])
  tf.debugging.assert_equal(collectedSteps[:, 1:2], collectedSteps[:, 2:3])
  return

# test that masked values aren't zeroed
def test_maskedValues():
  fake = _fake_AR(threshold=1e+5, scale=0.0)
  watcher = CSamplerWatcher(
    steps=10,
    tracked=dict(value=(32, 3)),
  )
  arg = dict(value=fake.x, model=fake.model, algorithmInterceptor=watcher.interceptor())
  _ = fake.interpolant.sample(**arg, index=0)
  
  collectedSteps = watcher.tracked('value')
  afterMask = collectedSteps[3:]
  beforeMask = collectedSteps[2]

  tf.debugging.assert_greater(3, watcher.iteration, 'Must collect 3 steps')
  for i in range(3, watcher.iteration + 1):
    tf.debugging.assert_equal(afterMask[i], beforeMask, 'Must be equal')
  return