def test_realtime_speed(dev_str, call):
if call in [helpers.np_call, helpers.jnp_call, helpers.mx_call]:
# convolutions not yet implemented in numpy or jax
# mxnet is unable to stack or expand zero-dimensional tensors
pytest.skip()
ivy.seed(0)
device = 'cpu'
batch_size = 1
num_timesteps = 1
num_cams = 1
num_feature_channels = 3
image_dims = [64, 64]
omni_img_dims = [90, 180]
esm = ESM(omni_image_dims=omni_img_dims, device=device)
memory = esm.empty_memory(batch_size, num_timesteps)
start_time = time.perf_counter()
for i in range(50):
obs = _get_dummy_obs(batch_size, num_timesteps, num_cams, image_dims, num_feature_channels, device)
memory = esm(obs, memory, batch_size=batch_size, num_timesteps=num_timesteps, num_cams=num_cams,
image_dims=image_dims)
memory_mean = memory.mean.numpy()
assert memory_mean.shape == tuple([batch_size, num_timesteps] + omni_img_dims + [3 + num_feature_channels])
assert memory_mean[0, 0, 0, 0, 0] == 0.
np.max(memory_mean)
end_time = time.perf_counter()
time_taken = end_time - start_time
assert time_taken < 20.
def test_slice_wrapped(self, dev_str, f, call, array_shape, num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
ivy.seed(0)
np.random.seed(0)
self._init(array_shape, num_processes)
assert len(self._dataset[-1:1].x) == 2
assert list(self._dataset[-1:1].x[0].shape) == array_shape
assert len(self._dataset[-1:1].x) == 2
assert list(self._dataset[-1:1].x[1].shape) == array_shape
assert len(self._dataset[9:11].x) == 2
assert list(self._dataset[9:11].x[0].shape) == array_shape
assert len(self._dataset[9:11].x) == 2
assert list(self._dataset[9:11].x[0].shape) == array_shape
check0 = not np.allclose(ivy.to_numpy(self._dataset[-1:1].x[0]),
ivy.to_numpy(self._x[8]))
check1 = not np.allclose(ivy.to_numpy(self._dataset[-1:1].x[1]),
ivy.to_numpy(self._x[0]))
check2 = not np.allclose(ivy.to_numpy(self._dataset[9:11].x[0]),
ivy.to_numpy(self._x[0]))
check3 = not np.allclose(ivy.to_numpy(self._dataset[9:11].x[1]),
ivy.to_numpy(self._x[1]))
assert check0 or check1 or check2 or check3
# close
self._dataset.close()
del self._dataset
def __init__(self, num_iters, compile_flag, interactive, dev_str, f):
# ivy
f = choose_random_framework() if f is None else f
ivy.set_framework(f)
ivy.seed(0)
# device
if dev_str is None:
dev_str = 'gpu:0' if ivy.gpu_is_available() else 'cpu'
self._dev_str = dev_str
# Load input images and poses
this_dir = os.path.dirname(os.path.realpath(__file__))
data = np.load(os.path.join(this_dir, 'nerf_data/tiny_nerf_data.npz'))
images = ivy.array(data['images'], 'float32', dev_str)
inv_ext_mats = ivy.array(data['poses'], 'float32', dev_str)
# intrinsics
focal_lengths = ivy.array(np.tile(np.reshape(data['focal'], (1, 1)), [100, 2]), 'float32', dev_str)
self._img_dims = images.shape[1:3]
pp_offsets = ivy.tile(ivy.array([[dim/2 - 0.5 for dim in self._img_dims]], dev_str=dev_str), [100, 1])
# train data
self._images = images[:100, ..., :3]
self._intrinsics = ivy_vision.focal_lengths_and_pp_offsets_to_intrinsics_object(
focal_lengths, pp_offsets, self._img_dims)
self._cam_geoms = ivy_vision.inv_ext_mat_and_intrinsics_to_cam_geometry_object(
inv_ext_mats[:100, 0:3], self._intrinsics)
# test data
self._test_img = images[101]
self._test_cam_geom = ivy_vision.inv_ext_mat_and_intrinsics_to_cam_geometry_object(
inv_ext_mats[101, 0:3], self._intrinsics.slice(0))
# train config
self._embed_length = 6
self._lr = 5e-4
self._num_samples = 64
self._num_iters = num_iters
# log config
self._interactive = interactive
self._log_freq = 1
self._vis_freq = 25 if self._interactive else -1
self._vis_log_dir = 'nerf_renderings'
if os.path.exists(self._vis_log_dir):
shutil.rmtree(self._vis_log_dir)
os.makedirs(self._vis_log_dir)
# model
self._model = Model(4, 256, self._embed_length, dev_str)
# compile
if compile_flag:
rays_o, rays_d = self._get_rays(self._cam_geoms.slice(0))
target = self._images[0]
self._loss_fn = ivy.compile_fn(self._loss_fn, False,
example_inputs=[self._model, rays_o, rays_d, target, self._model.v])
def test_seed(seed_val, dtype_str, tensor_fn, dev_str, call):
# smoke test
ivy.seed(seed_val)
# compilation test
if call in [helpers.torch_call]:
# pytorch scripting does not support functions with None return
return
helpers.assert_compilable(ivy.seed)
def test_shuffle(x, dtype_str, tensor_fn, dev_str, call):
# smoke test
x = tensor_fn(x, dtype_str, dev_str)
ret = ivy.shuffle(x)
# type test
assert ivy.is_array(ret)
# cardinality test
assert ret.shape == x.shape
# value test
ivy.seed(0)
first_shuffle = call(ivy.shuffle, x)
ivy.seed(0)
second_shuffle = call(ivy.shuffle, x)
assert np.array_equal(first_shuffle, second_shuffle)
# compilation test
helpers.assert_compilable(ivy.shuffle)
def shuffle(self, seed_value=None):
"""
Shuffle entries in all sub-arrays, such that they are still aligned along axis 0.
:param seed_value: random seed to use for array shuffling
:type seed_value: int
"""
return_dict = dict()
if seed_value is None:
seed_value = _random.randint(0, 1000)
for key, value in sorted(self.items()):
if isinstance(value, Container):
return_dict[key] = value.shuffle(seed_value)
else:
_ivy.seed(seed_value)
return_dict[key] = _ivy.shuffle(value)
return Container(return_dict)
def main(env_str, steps=100, iters=10000, lr=0.001, seed=0, log_freq=100, vis_freq=1000, visualize=True, f=None):
# config
f = choose_random_framework(excluded=['numpy']) if f is None else f
ivy.seed(seed)
env = getattr(ivy_gym, env_str)()
starting_obs = env.reset()
# policy
in_size = starting_obs.shape[0]
ac_dim = env.action_space.shape[0]
policy = Policy(in_size, ac_dim)
# compile loss function
compiled_loss_fn = ivy.compile_fn(lambda initial_state, pol_vs:
loss_fn(env, initial_state, policy, pol_vs, steps),
False, example_inputs=[env.get_state(), policy.v])
# optimizer
optimizer = ivy.Adam(lr=lr)
# train
scores = []
for iteration in range(iters):
if iteration % vis_freq == 0 and visualize:
obs = env.reset()
env.render()
for _ in range(steps):
ac = policy(obs)
obs, _, _, _ = env.step(ac)
env.render()
env.reset()
if iteration == 0:
print('\nCompiling loss function for {} environment steps... This may take a while...\n'.format(steps))
score = train_step(compiled_loss_fn, optimizer, env.get_state(), policy, f)
if iteration == 0:
print('\nLoss function compiled!\n')
print('iteration {} score {}'.format(iteration, ivy.to_numpy(score).item()))
scores.append(ivy.to_numpy(score)[0])
if len(scores) == log_freq:
print('\nIterations: {} Mean Score: {}\n'.format(iteration + 1, np.mean(scores)))
scores.clear()
def __init__(self, controller, controller_proj, output_proj, output_dim, ctrl_input_size, ctrl_output_size,
total_parameter_num, memory_size, memory_vector_dim, read_head_num, write_head_num, v=None, usage=None,
addressing_mode='content_and_location', shift_range=1, clip_value=20, init_value=1e-6,
sequential_writing=False, retroactive_updates=False, retroactive_discount=0.96, with_erase=True,
seed=0):
# vanilla ntm
self._memory_size = memory_size
self._memory_vector_dim = memory_vector_dim
self._read_head_num = read_head_num
self._write_head_num = write_head_num
self._init_value = init_value
self._addressing_mode = addressing_mode
self._clip_value = clip_value
self._output_dim = output_dim
self._shift_range = shift_range
self._num_parameters_per_head = self._memory_vector_dim + 1 + 1 + (self._shift_range * 2 + 1) + 1
self._num_heads = self._read_head_num + self._write_head_num
ivy.seed(seed)
# fns + classes
self._controller = controller
self._controller_proj = controller_proj
self._output_proj = output_proj
# usage
if usage is not None:
self._usage = usage
else:
self._usage = ivy.zeros([memory_size, ])
# step
self._step = 0
# MERLIN changes
self._sequential_writing = sequential_writing
self._retroactive_updates = retroactive_updates
self._retroactive_discount = retroactive_discount
self._with_erase = with_erase
# variables
ivy.Module.__init__(self, 'cpu')
def main(seed=0, compile_mode=False, dev_strs=None):
ivy.seed(seed)
data_loader_spec_args = {
'batch_size': 2,
'dev_strs': [ivy.default(lambda: dev_strs[0], None, True)]
}
trainer_spec_args = {
'total_iterations': 10,
'ld_chkpt': False,
'log_freq': 1,
'initial_learning_rate': 0.1,
'compile_mode': compile_mode,
'dev_strs': dev_strs
}
trainer = trainer_builder.build_trainer(
ExampleDataLoader,
ExampleNetwork,
ExampleTrainer,
data_loader_spec_args=data_loader_spec_args,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
trainer.close()
def test_single(self, dev_str, f, call, array_shape, num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
ivy.seed(0)
np.random.seed(0)
self._init(array_shape, num_processes)
assert list(self._dataset[0].x.shape) == array_shape
assert list(self._dataset[4].x.shape) == array_shape
assert list(self._dataset[8].x.shape) == array_shape
check0 = not np.allclose(ivy.to_numpy(self._dataset[0].x),
ivy.to_numpy(self._x[0]))
check1 = not np.allclose(ivy.to_numpy(self._dataset[1].x),
ivy.to_numpy(self._x[1]))
check2 = not np.allclose(ivy.to_numpy(self._dataset[2].x),
ivy.to_numpy(self._x[2]))
check3 = not np.allclose(ivy.to_numpy(self._dataset[3].x),
ivy.to_numpy(self._x[3]))
check4 = not np.allclose(ivy.to_numpy(self._dataset[4].x),
ivy.to_numpy(self._x[4]))
check5 = not np.allclose(ivy.to_numpy(self._dataset[5].x),
ivy.to_numpy(self._x[5]))
check6 = not np.allclose(ivy.to_numpy(self._dataset[6].x),
ivy.to_numpy(self._x[6]))
check7 = not np.allclose(ivy.to_numpy(self._dataset[7].x),
ivy.to_numpy(self._x[7]))
check8 = not np.allclose(ivy.to_numpy(self._dataset[8].x),
ivy.to_numpy(self._x[8]))
assert check0 or check1 or check2 or check3 or check4 or check5 or check6 or check7 or check8
# close
self._dataset.close()
del self._dataset
def test_reduced_cost_after_checkpoint_load(dev_str, call, compile_mode):
if call is helpers.np_call:
# numpy does not support gradients, required for training
pytest.skip()
if call is helpers.jnp_call and ivy.wrapped_mode():
# Jax does not support ivy.Array instances when calling _jax.grad()
pytest.skip()
example_dir = os.path.relpath(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'../ivy_builder_demos'))
# currently only PyTorch supports graph compilation
compile_mode = compile_mode if ivy.current_framework_str(
) == 'torch' else False
# dataset dirs specification
dataset_dirs_args = dict()
# dataset specification
dataset_spec_filepath = os.path.join(example_dir, 'json_specs',
'dataset_spec.json.example')
dataset_spec_args = builder.parse_json_to_cont(dataset_spec_filepath)
# data loader specification
data_loader_spec_filepath = os.path.join(example_dir, 'json_specs',
'data_loader_spec.json.example')
data_loader_spec_args = builder.parse_json_to_cont(
data_loader_spec_filepath)
# network specification
network_spec_filepath = os.path.join(example_dir, 'json_specs',
'network_spec.json.example')
network_spec_args = builder.parse_json_to_cont(network_spec_filepath)
builder_helpers.remove_dirs()
ivy.seed(0)
trainer_spec_args = {
'total_iterations': 1,
'ld_chkpt': False,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoader,
ExampleNetwork,
ExampleTrainer,
dataset_dirs_args=dataset_dirs_args,
dataset_dirs_class=ExampleDatasetDirs,
dataset_spec_args=dataset_spec_args,
dataset_spec_class=ExampleDatasetSpec,
data_loader_spec_args=data_loader_spec_args,
data_loader_spec_class=ExampleDataLoaderSpec,
network_spec_args=network_spec_args,
network_spec_class=ExampleNetworkSpec,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
initial_cost = trainer._total_cost
assert trainer._global_step == 1
trainer.close()
ivy.seed(0)
steps_to_take_first = 10
trainer_spec_args = {
'total_iterations': steps_to_take_first,
'ld_chkpt': False,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoader,
ExampleNetwork,
ExampleTrainer,
dataset_dirs_args=dataset_dirs_args,
dataset_dirs_class=ExampleDatasetDirs,
dataset_spec_args=dataset_spec_args,
dataset_spec_class=ExampleDatasetSpec,
data_loader_spec_args=data_loader_spec_args,
data_loader_spec_class=ExampleDataLoaderSpec,
network_spec_args=network_spec_args,
network_spec_class=ExampleNetworkSpec,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
ten_step_cost = trainer._total_cost
assert trainer._global_step == steps_to_take_first
trainer.close()
assert initial_cost > ten_step_cost
steps_to_take_second = 20
trainer_spec_args = {
'total_iterations': steps_to_take_second,
'ld_chkpt': True,
'save_freq': 1,
'compile_mode': compile_mode
}
trainer = builder.build_trainer(
ExampleDataLoader,
ExampleNetwork,
ExampleTrainer,
dataset_dirs_args=dataset_dirs_args,
dataset_dirs_class=ExampleDatasetDirs,
dataset_spec_args=dataset_spec_args,
dataset_spec_class=ExampleDatasetSpec,
data_loader_spec_args=data_loader_spec_args,
data_loader_spec_class=ExampleDataLoaderSpec,
network_spec_args=network_spec_args,
network_spec_class=ExampleNetworkSpec,
trainer_spec_args=trainer_spec_args)
trainer.setup()
trainer.train()
twenty_step_cost = trainer._total_cost
assert trainer._global_step == steps_to_take_second
trainer.close()
assert ten_step_cost > twenty_step_cost
builder_helpers.remove_dirs()
def __init__(self, spec: TrainerSpec) -> None:
# specification
self._spec = spec
# uninitialized variables
if spec.starting_iteration is not None:
self._starting_iteration = spec.starting_iteration
else:
self._starting_iteration = 0
self._total_iterations = None
# trainer variables
self._global_step = 0
self._moving_average_loss = 0
# set seed
np.random.seed(self._spec.seed)
ivy.seed(self._spec.seed)
# uninitialized variables
self._chkpt = None
self._chkpt_manager = None
# summary writer
try:
from torch.utils.tensorboard import SummaryWriter
except ModuleNotFoundError:
SummaryWriter = None
if SummaryWriter is not None:
self._writer = SummaryWriter(os.path.join(self._spec.log_dir, 'tnsrbrd'))
else:
self._writer = None
# profiler
self._profiling = self._spec.steps_to_profile > 0
if self._profiling:
self._profiler = ivy.Profiler(self._spec.log_dir)
else:
self._profiler = None
# timing
self._start_time = time.perf_counter()
# batch
self._training_batch = None
# network
self._network = self._spec.network
self._net_spec = self._network.spec
self._partial_grad_updates = bool(self._net_spec.v_keychains)
# multi-dev
self._dev_str = ivy.default(lambda: self._spec.dev_strs[0], ivy.default_device(), True)
if len(self._spec.dev_strs) > 1:
if self._network.built:
raise Exception('Network must use either explicit or on_call build modes if training on multiple'
'devices, but the network was already built using on_init method.')
ret_fn = lambda ret: ivy.unify_iter(ret, self._spec.dev_strs[0], 'mean', transpose=True)
dev_mapper = ivy.DevMapperMultiProc(
self.__getattribute__(self._spec.dev_map_fn), ret_fn, self._spec.dev_strs,
constant={'network': self._network})
self._multi_dev = True
else:
dev_mapper = None
self._multi_dev = False
# device manager
if (self._multi_dev and self._spec.tune_device_allocation) or self._spec.tune_splitting:
self._dev_manager = ivy.DevManager(
dev_mapper, self._spec.dev_strs, tune_dev_alloc=(self._multi_dev and self._spec.tune_device_allocation),
tune_dev_splits=self._spec.tune_splitting)
else:
self._dev_manager = None
# compilation
self._compile_network_once_tuned = False
self._compile_optimizer_once_tuned = False
