#include <string>

#include <iostream>

#include <CoverageControlTorch/cnn_module.h>

using namespace CoverageControlTorch;

void LoadDataset(std::string const &data_dir, int const image_size, int const output_dim, torch::Tensor &maps, torch::Tensor &features, torch::Tensor &features_mean, torch::Tensor &features_std) {

std::string local_maps_file = data_dir + "/local_maps.pt";

std::string comm_maps_file = data_dir + "/comm_maps.pt";

std::string obstacle_maps_file = data_dir + "/obstacle_maps.pt";

std::string features_file = data_dir + "/normalized_coverage_features.pt";

torch::Tensor local_maps;

torch::load(local_maps, local_maps_file);

local_maps = local_maps.unsqueeze(2).view({-1, 1, image_size, image_size});

torch::Tensor comm_maps;

torch::load(comm_maps, comm_maps_file);

comm_maps = comm_maps.to_dense().view({-1, 2, image_size, image_size});

torch::Tensor obstacle_maps;

torch::load(obstacle_maps, obstacle_maps_file);

obstacle_maps = obstacle_maps.to_dense().unsqueeze(2).view({-1, 1, image_size, image_size});

torch::load(features, features_file);

features = features.view({-1, features.size(2)});

features = features.index({Slice(), Slice(0, output_dim)}).to(torch::kCPU);

maps = torch::cat({local_maps, comm_maps, obstacle_maps}, 1).to(torch::kCPU);

torch::load(features_mean, data_dir + "/coverage_features_mean.pt");

torch::load(features_std, data_dir + "/coverage_features_std.pt");

std::cout << "maps shape: " << maps.sizes() << std::endl;

}

int main(int argc, char* argv[]) {

if (argc < 2) {

std::cout << "Usage: ./test_cnn <yaml>" << std::endl;

return 1;

}

std::string config_file = std::string(argv[1]);

torch::Device device(torch::kCPU);

if (torch::cuda::is_available()) {

device = torch::kCUDA;

std::cout << "Using CUDA" << std::endl;

}

YAML::Node config = YAML::LoadFile(config_file);

auto cnn_config = config["CNN"];

std::string data_dir = config["pDataDir"].as<std::string>();

CoverageControlCNN model(cnn_config);

torch::load(model, data_dir + config["CNNTraining"]["ModelCkpt"].as<std::string>());

model->to(device);

int image_size = cnn_config["ImageSize"].as<int>();

int output_dim = cnn_config["OutputDim"].as<int>();

torch::Tensor maps, features, features_mean, features_std;

LoadDataset(data_dir, image_size, output_dim, maps, features, features_mean, features_std);

std::cout << "Loaded dataset" << std::endl;

maps = maps.to(torch::kCPU);

torch::Tensor sub_maps = maps.index({Slice(0, 1000)});

sub_maps = sub_maps.to(device);

torch::Tensor sub_features = features.index({Slice(0, 1000)});

std::cout << "submaps created" << std::endl;

auto pred = model->forward(sub_maps).to(torch::kCPU);

std::cout << "Pred done" << std::endl;

auto loss = torch::mse_loss(pred, sub_features);

std::cout << "Val loss: " << loss << std::endl;

return 0;

}

#ifndef COVERAGECONTROL_CNN_MODULE_H_

#define COVERAGECONTROL_CNN_MODULE_H_

#include <thread>

#include <string>

#include <yaml-cpp/yaml.h>

#include <torch/torch.h>

#include "cnn_backbone.h"

using namespace torch::indexing;

namespace F = torch::nn::functional;

namespace CoverageControlTorch {

struct CoverageControlCNNImpl : torch::nn::Module {

int input_dim_ = 4;

int output_dim_ = 7;

int num_layers_ = 2;

int latent_size_ = 8;

int kernel_size_ = 3;

int image_size_ = 32;

CNNBackbone cnn_backbone_;

torch::nn::Linear linear_;

CoverageControlCNNImpl(YAML::Node config) : CoverageControlCNNImpl(

config["InputDim"].as<int>(),

config["OutputDim"].as<int>(),

config["NumLayers"].as<int>(),

config["LatentSize"].as<int>(),

config["KernelSize"].as<int>(),

config["ImageSize"].as<int>()) {}

CoverageControlCNNImpl(int input_dim, int output_dim, int num_layers, int latent_size, int kernel_size, int image_size) :

input_dim_(input_dim),

output_dim_(output_dim),

num_layers_(num_layers),

latent_size_(latent_size),

kernel_size_(kernel_size),

image_size_(image_size),

cnn_backbone_(register_module("cnn_backbone", CNNBackbone(input_dim_, output_dim_, num_layers_, latent_size_, kernel_size_, image_size_))),

linear_(register_module("linear", torch::nn::Linear(2 * output_dim_, output_dim))) {

}

torch::Tensor forward(torch::Tensor x) {

x = cnn_backbone_->forward(x);

x = linear_->forward(x);

return x;

}

};

TORCH_MODULE(CoverageControlCNN);

} // CoverageControlTorch

#endif // COVERAGECONTROL_CNN_MODULE_H_

#include "cnn_module.h"

CoverageControlCNN model(cnn_config_);

cnn_config_ = config_["CNN"];