Skip to content

Commit 3d74386

Browse files
Semyon1104Semyon1104
committed
fix json for yolo, fix split, check concats
1 parent ec525fa commit 3d74386

File tree

5 files changed

+110
-52
lines changed

5 files changed

+110
-52
lines changed

‎app/Converters/parser_onnx.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ def convert_pt_to_onnx(pt_model_path, onnx_model_path=None):
1111
onnx_model_path = pt_model_path.replace('.pt', '.onnx')
1212

1313
model = YOLO(pt_model_path)
14-
model.export(format="onnx", dynamic=False, simplify=True)
14+
model.export(format="onnx", dynamic=False, simplify=False)
1515

1616
return onnx_model_path
1717

@@ -156,7 +156,7 @@ def default(self, obj):
156156

157157
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
158158

159-
MODEL_PATH = os.path.join(BASE_DIR, 'docs\\models', 'densenet121_Opset16.onnx')
160-
MODEL_DATA_PATH = os.path.join(BASE_DIR, 'docs\\jsons', 'densenet121_Opset16_onnx_model.json')
159+
MODEL_PATH = os.path.join(BASE_DIR, 'docs\\models', 'yolo11x-cls.pt')
160+
MODEL_DATA_PATH = os.path.join(BASE_DIR, 'docs\\jsons', 'yolo11x-cls_onnx_model.json')
161161

162162
onnx_to_json(MODEL_PATH, MODEL_DATA_PATH)

‎app/Graph/build.cpp‎

Lines changed: 97 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
#include <regex>
44
#include <set>
55
#include <unordered_map>
6+
#include <unordered_set>
67

78
void build_graph_linear(it_lab_ai::Tensor& input, it_lab_ai::Tensor& output,
89
const std::string& json_path, bool comments,
@@ -288,6 +289,10 @@ void build_graph(it_lab_ai::Tensor& input, it_lab_ai::Tensor& output,
288289
it_lab_ai::ImplType impl1 = parallel ? it_lab_ai::kTBB : it_lab_ai::kDefault;
289290
it_lab_ai::ImplType impl2 = parallel ? it_lab_ai::kSTL : it_lab_ai::kDefault;
290291

292+
std::unordered_map<std::string, std::vector<std::string>> concat_connections;
293+
std::unordered_map<std::string, std::vector<int>> concat_orders;
294+
std::unordered_map<std::string, std::unordered_set<std::string>> concat_connected_inputs;
295+
291296
std::unordered_map<std::string, std::vector<int64_t>> layer_parameters;
292297
std::unordered_map<std::string, float> float_parameters;
293298
std::string last_constant_name;
@@ -557,17 +562,39 @@ void build_graph(it_lab_ai::Tensor& input, it_lab_ai::Tensor& output,
557562
if (layer_data["attributes"].contains("axis")) {
558563
axis = layer_data["attributes"]["axis"];
559564
}
565+
if (layer_data.contains("inputs")) {
566+
for (const auto& input_name : layer_data["inputs"]) {
567+
std::string input_tensor = input_name.get<std::string>();
568+
std::string base_input_name = get_base_layer_name(input_tensor);
569+
concat_connections[layer_name].push_back(base_input_name);
570+
}
571+
}
560572
auto concat_layer = std::make_shared<it_lab_ai::ConcatLayer>(axis);
561573
concat_layer->setName(it_lab_ai::kConcat);
562574
layer = concat_layer;
575+
concat_connected_inputs[layer_name] = std::unordered_set<std::string>();
563576
} else if (layer_type == "Split") {
564577
int axis = 0;
565-
std::vector<int> splits;
578+
std::vector<int64_t> splits;
566579
size_t num_outputs = 2;
567580

568581
if (layer_data["attributes"].contains("axis")) {
569582
axis = layer_data["attributes"]["axis"];
570583
}
584+
if (layer_data.contains("inputs") && layer_data["inputs"].is_array()) {
585+
auto inputs = layer_data["inputs"];
586+
if (inputs.size() >= 2) {
587+
std::string constant_name = inputs[1].get<std::string>();
588+
constant_name = get_base_layer_name(constant_name);
589+
590+
if (layer_parameters.count(constant_name)) {
591+
splits = layer_parameters[constant_name];
592+
} else if (constant_name.find("onnx::") != constant_name.npos) {
593+
splits = last_constant_value;
594+
layer_parameters[constant_name] = last_constant_value;
595+
}
596+
}
597+
}
571598
if (layer_data.contains("weights") &&
572599
layer_data["weights"].is_array()) {
573600
for (const auto& s : layer_data["weights"]) {
@@ -642,6 +669,10 @@ void build_graph(it_lab_ai::Tensor& input, it_lab_ai::Tensor& output,
642669
std::make_shared<it_lab_ai::EWLayer>(ew_operation, value, 0.0f);
643670
ew_layer->setName(it_lab_ai::kElementWise);
644671
layer = ew_layer;
672+
/*if (comments) {
673+
std::cout << "Created binary " << layer_type << " operation with "
674+
<< value <<"scalar" << std::endl;
675+
}*/
645676
} else if (layer_type == "Add") {
646677
ew_operation = "linear";
647678
auto ew_layer =
@@ -676,11 +707,7 @@ void build_graph(it_lab_ai::Tensor& input, it_lab_ai::Tensor& output,
676707
auto bin_layer = std::make_shared<it_lab_ai::BinaryOpLayer>(op);
677708
bin_layer->setName(it_lab_ai::kBinaryOp);
678709
layer = bin_layer;
679-
680-
if (comments) {
681-
std::cout << "Created binary " << layer_type
682-
<< " operation with tensor inputs" << std::endl;
683-
}
710+
684711
}
685712
} else if (layer_type == "Gemm") {
686713
it_lab_ai::Tensor tensor = it_lab_ai::create_tensor_from_json(
@@ -1019,6 +1046,8 @@ void build_graph(it_lab_ai::Tensor& input, it_lab_ai::Tensor& output,
10191046
for (const auto& input_name : layer_data["inputs"]) {
10201047
std::string input_tensor = input_name.get<std::string>();
10211048

1049+
1050+
10221051
// Проверяем, является ли вход выходом сплит-слоя
10231052
std::regex split_output_pattern("(.+)_output_(\\d+)$");
10241053
std::smatch matches;
@@ -1206,50 +1235,78 @@ void build_graph(it_lab_ai::Tensor& input, it_lab_ai::Tensor& output,
12061235
/*if (comments) {
12071236
std::cout << "\n=== ESTABLISHING CONNECTIONS ===" << std::endl;
12081237
}*/
1238+
std::vector<int> order = {};
12091239

12101240
for (const auto& [source_name, target_name] : connection_list) {
1211-
// Убираем проверку на сплит-выходы - они тоже должны быть подключены
1212-
12131241
if (name_to_layer.count(source_name) && name_to_layer.count(target_name)) {
1242+
// Обработка Concat слоев
1243+
if (target_name.find("Concat") != std::string::npos ||
1244+
name_to_layer[target_name]->getName() == it_lab_ai::kConcat) {
1245+
// Проверяем, есть ли этот concat в нашем списке
1246+
if (concat_connections.find(target_name) != concat_connections.end()) {
1247+
// Находим индекс этого источника в ожидаемых входах concat
1248+
const auto& expected_inputs = concat_connections[target_name];
1249+
auto it = std::find(expected_inputs.begin(), expected_inputs.end(),
1250+
source_name);
1251+
1252+
if (it != expected_inputs.end()) {
1253+
int input_index = static_cast<int>(std::distance(expected_inputs.begin(), it));
1254+
1255+
// Добавляем индекс в порядок для этого concat
1256+
concat_orders[target_name].push_back(input_index);
1257+
1258+
// Отмечаем, что этот вход подключен
1259+
concat_connected_inputs[target_name].insert(source_name);
1260+
1261+
if (comments) {
1262+
std::cout << "Concat connection: " << source_name << " -> "
1263+
<< target_name << " (index: " << input_index << ")"
1264+
<< std::endl;
1265+
}
1266+
1267+
// Проверяем, все ли входы подключены
1268+
if (concat_connected_inputs[target_name].size() ==
1269+
concat_connections[target_name].size()) {
1270+
// Все входы подключены - устанавливаем порядок
1271+
auto concat_layer =
1272+
std::dynamic_pointer_cast<it_lab_ai::ConcatLayer>(
1273+
name_to_layer[target_name]);
1274+
if (concat_layer) {
1275+
concat_layer->setInputOrder(concat_orders[target_name]);
1276+
1277+
if (comments) {
1278+
std::cout
1279+
<< "=== ALL INPUTS CONNECTED TO CONCAT: " << target_name
1280+
<< " ===" << std::endl;
1281+
std::cout << "Expected inputs: ";
1282+
for (const auto& inp : concat_connections[target_name]) {
1283+
std::cout << inp << " ";
1284+
}
1285+
std::cout << std::endl;
1286+
1287+
std::cout << "Actual order: ";
1288+
for (size_t i = 0; i < concat_orders[target_name].size();
1289+
++i) {
1290+
std::cout << concat_orders[target_name][i];
1291+
if (i < concat_orders[target_name].size() - 1)
1292+
std::cout << ", ";
1293+
}
1294+
std::cout << std::endl;
1295+
}
1296+
}
1297+
}
1298+
}
1299+
}
1300+
}
1301+
12141302
try {
1215-
//if (comments) {
1216-
// std::cout << "Connecting: " << source_name << " -> " << target_name;
1217-
// std::cout << " (ID: " << name_to_layer[source_name]->getID()
1218-
// << " -> ID: " << name_to_layer[target_name]->getID() << ")"
1219-
// << std::endl;
1220-
1221-
// // Дополнительная информация для сплит-соединений
1222-
// std::regex split_output_pattern("(.+)_output_(\\d+)$");
1223-
// std::smatch matches;
1224-
// if (std::regex_search(source_name, matches, split_output_pattern)) {
1225-
// std::string split_layer_name = matches[1].str();
1226-
// int output_index = std::stoi(matches[2].str());
1227-
// std::cout << " [SPLIT] Output index: " << output_index
1228-
// << std::endl;
1229-
// }
1230-
//}
12311303
graph.makeConnection(*name_to_layer[source_name],
12321304
*name_to_layer[target_name]);
1233-
/*if (comments) {
1234-
std::cout << " Success" << std::endl;
1235-
}*/
1305+
12361306
} catch (const std::exception& e) {
12371307
std::cerr << "Failed: " << source_name << " -> " << target_name << " : "
12381308
<< e.what() << std::endl;
12391309
}
1240-
} else {
1241-
/*if (comments) {
1242-
std::cerr << "Warning: Missing layer for connection " << source_name
1243-
<< " -> " << target_name << std::endl;
1244-
if (!name_to_layer.count(source_name)) {
1245-
std::cerr << " Source layer '" << source_name << "' not found"
1246-
<< std::endl;
1247-
}
1248-
if (!name_to_layer.count(target_name)) {
1249-
std::cerr << " Target layer '" << target_name << "' not found"
1250-
<< std::endl;
1251-
}
1252-
}*/
12531310
}
12541311
}
12551312
for (auto& split_dist : split_distribution) {

‎include/layers/ConcatLayer.hpp‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ class ConcatLayer : public Layer {
1515

1616
void run(const std::vector<Tensor>& input,
1717
std::vector<Tensor>& output) override;
18-
18+
void setInputOrder(const std::vector<int>& order) { input_order_ = order; }
1919
static std::string get_name() { return "ConcatLayer"; }
2020

2121
#ifdef ENABLE_STATISTIC_WEIGHTS
@@ -24,14 +24,15 @@ class ConcatLayer : public Layer {
2424

2525
private:
2626
int64_t axis_;
27-
27+
std::vector<int> input_order_;
2828
void validate_inputs(const std::vector<Tensor>& inputs) const;
2929
int64_t normalize_axis(size_t rank) const;
3030
Shape calculate_output_shape(const std::vector<Tensor>& inputs) const;
31-
31+
std::vector<Tensor> reorderInputs(const std::vector<Tensor>& inputs) const;
3232
template <typename T>
3333
void concatenate(const std::vector<Tensor>& inputs,
3434
Tensor& output) const {
35+
std::vector<Tensor> ordered_inputs = reorderInputs(inputs);
3536
Shape output_shape = calculate_output_shape(inputs);
3637
std::vector<T> output_data(output_shape.count(), 0);
3738

‎include/layers/SplitLayer.hpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ namespace it_lab_ai {
1010

1111
class SplitLayer : public Layer {
1212
public:
13-
SplitLayer(int axis, std::vector<int> splits)
13+
SplitLayer(int axis, std::vector<int64_t> splits)
1414
: axis_(axis), splits_(std::move(splits)) {}
1515

1616
SplitLayer(int axis, int num_outputs)
@@ -26,7 +26,7 @@ class SplitLayer : public Layer {
2626

2727
private:
2828
int axis_;
29-
std::optional<std::vector<int>> splits_;
29+
std::optional<std::vector<int64_t>> splits_;
3030
std::optional<int> num_outputs_;
3131

3232
void validate(const Tensor& input) const;

‎src/layers/SplitLayer.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ void SplitLayer::split_impl(const Tensor& input,
3232
const Shape& shape = input.get_shape();
3333
const int axis = get_normalized_axis(static_cast<int>(shape.dims()));
3434

35-
std::vector<int> part_sizes;
35+
std::vector<int64_t> part_sizes;
3636
if (splits_) {
3737
part_sizes = *splits_;
3838
} else {
@@ -41,7 +41,7 @@ void SplitLayer::split_impl(const Tensor& input,
4141
const int remainder = total_size % *num_outputs_;
4242

4343
part_sizes.reserve(*num_outputs_);
44-
for (int i = 0; i < *num_outputs_; ++i) {
44+
for (int64_t i = 0; i < *num_outputs_; ++i) {
4545
part_sizes.push_back(i < remainder ? base_size + 1 : base_size);
4646
}
4747
}
@@ -99,8 +99,8 @@ void SplitLayer::validate(const Tensor& input) const {
9999
const int axis_size = static_cast<int>(input.get_shape()[axis]);
100100

101101
if (splits_) {
102-
int sum = 0;
103-
for (int s : *splits_) {
102+
int64_t sum = 0;
103+
for (int64_t s : *splits_) {
104104
if (s <= 0) throw std::runtime_error("Split size must be positive");
105105
sum += s;
106106
}

0 commit comments

Comments
 (0)