TestTorchTableFunctions.h File Reference

#include "QueryEngine/heavydbTypes.h"

Include dependency graph for TestTorchTableFunctions.h:

This graph shows which files directly or indirectly include this file:

Go to the source code of this file.

Functions
EXTENSION_NOINLINE int32_t	tf_test_runtime_torch (TableFunctionManager &mgr, Column< int64_t > &input, Column< int64_t > &output)
template<typename T >
TEMPLATE_NOINLINE int32_t	tf_test_runtime_torch_template__template (TableFunctionManager &mgr, const Column< T > &input, Column< T > &output)
EXTENSION_NOINLINE int32_t	tf_test_torch_regression (TableFunctionManager &mgr, const ColumnList< double > &features, int32_t batch_size, bool use_gpu, bool save_model, const TextEncodingNone &model_filename, Column< double > &output)
EXTENSION_NOINLINE int32_t	tf_test_torch_generate_random_column (TableFunctionManager &mgr, int32_t num_elements, Column< double > &output)
EXTENSION_NOINLINE int32_t	tf_test_torch_load_model (TableFunctionManager &mgr, const TextEncodingNone &model_filename, Column< bool > &output)

Function Documentation

EXTENSION_NOINLINE int32_t tf_test_runtime_torch	(	TableFunctionManager &	mgr,
		Column< int64_t > &	input,
		Column< int64_t > &	output
	)

Definition at line 43 of file TestTorchTableFunctions.cpp.

                                                       {
  return 0;
}

template<typename T >

TEMPLATE_NOINLINE int32_t tf_test_runtime_torch_template__template	(	TableFunctionManager &	mgr,
		const Column< T > &	input,
		Column< T > &	output
	)

Definition at line 51 of file TestTorchTableFunctions.cpp.

                                                            {
  return 0;
}

EXTENSION_NOINLINE int32_t tf_test_torch_generate_random_column	(	TableFunctionManager &	mgr,
		int32_t	num_elements,
		Column< double > &	output
	)

Definition at line 70 of file TestTorchTableFunctions.cpp.

References TableFunctionManager::set_output_row_size().

                                                                                        {
  mgr.set_output_row_size(num_elements);
  torch::Tensor random = torch::randn({num_elements}, at::dtype(at::kDouble));
  random = random.unsqueeze(1);
  double* data_ptr = (double*)random.data_ptr();

  for (int32_t i = 0; i < num_elements; ++i) {
    output[i] = *data_ptr++;
  }

  return num_elements;
}

Here is the call graph for this function:

EXTENSION_NOINLINE int32_t tf_test_torch_load_model	(	TableFunctionManager &	mgr,
		const TextEncodingNone &	model_filename,
		Column< bool > &	output
	)

Definition at line 226 of file TestTorchTableFunctions.cpp.

References TextEncodingNone::getString(), boost::serialization::load(), and TableFunctionManager::set_output_row_size().

                                               {
  mgr.set_output_row_size(1);
  torch::jit::script::Module module;
  try {
    module = torch::jit::load(model_filename.getString());
  } catch (const std::exception& e) {
    return mgr.ERROR_MESSAGE("Error loading torchscript model: " + e.what());
  }

  output[0] = true;
  return 1;
}

Here is the call graph for this function:

EXTENSION_NOINLINE int32_t tf_test_torch_regression	(	TableFunctionManager &	mgr,
		const ColumnList< double > &	features,
		int32_t	batch_size,
		bool	use_gpu,
		bool	save_model,
		const TextEncodingNone &	model_filename,
		Column< double > &	output
	)

Definition at line 146 of file TestTorchTableFunctions.cpp.

References _test_torch_tfs_device, f(), get_batch(), TextEncodingNone::getString(), ColumnList< T >::numCols(), poly_desc(), boost::serialization::save(), and TableFunctionManager::set_output_row_size().

                                                 {
  int32_t poly_degree = features.numCols();
  // we output target and trained coefficients + bias
  int32_t output_size = (poly_degree + 1) * 2;
  mgr.set_output_row_size(output_size);
  std::srand(std::time(nullptr));  // not ideal RNG, but fine for test purpooses
#ifdef HAVE_CUDA_TORCH
  if (torch::cuda::is_available() && use_gpu) {
    _test_torch_tfs_device = torch::kCUDA;
  }
#endif

  auto W_target = torch::randn({poly_degree, 1}, at::device(_test_torch_tfs_device)) * 5;
  auto b_target = torch::randn({1}, at::device(_test_torch_tfs_device)) * 5;

  // Define the model and optimizer
  auto fc = torch::nn::Linear(W_target.size(0), 1);
  fc->to(_test_torch_tfs_device);
  torch::optim::SGD optim(fc->parameters(), .1);

  float loss = 0;
  int64_t batch_idx = 0;

  while (++batch_idx) {
    // Get data
    torch::Tensor batch_x, batch_y;
    std::tie(batch_x, batch_y) = get_batch(features, W_target, b_target, batch_size);

    // Reset gradients
    optim.zero_grad();

    // Forward pass
    auto output = torch::smooth_l1_loss(fc(batch_x), batch_y);
    loss = output.item<float>();

    // Backward pass
    output.backward();

    // Apply gradients
    optim.step();

    // Stop criterion
    if (loss < 1e-3f)
      break;
  }

  if (save_model) {
    torch::save(fc, model_filename.getString());
  }

  // output column with target + trained coefficients ordered by degree, then bias
  torch::Tensor output_coefficients = fc->weight.view({-1}).cpu();
  torch::Tensor goal_coefficients = W_target.view({-1}).cpu();
  int32_t out_column_idx, input_idx;
  for (out_column_idx = 0, input_idx = 0; input_idx < output_coefficients.size(0);
       ++input_idx) {
    output[out_column_idx++] = output_coefficients[input_idx].item<float>();
    output[out_column_idx++] = goal_coefficients[input_idx].item<float>();
  }
  output[out_column_idx++] = fc->bias[0].item<float>();
  output[out_column_idx] = b_target[0].item<float>();

  std::fprintf(stdout, "Loss: %lf after %ld batches\n", loss, batch_idx);
  std::fprintf(stdout,
               "==> Learned function:\t%s\n",
               poly_desc(output_coefficients, fc->bias).c_str());
  std::fprintf(stdout,
               "==> Actual function:\t%s\n",
               poly_desc(W_target.view({-1}).cpu(), b_target).c_str());

  return output_size;
}

Here is the call graph for this function: