Skip to content

Commit 83f81b8

Browse files
nihuinihui
committed
w
1 parent aca6df2 commit 83f81b8

41 files changed

Lines changed: 782 additions & 244 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎tools/pnnx/src/pass_ncnn/F_glu.cpp‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,16 +36,17 @@ pnnx.Output output 1 0 out
3636
const int batch_index = op->inputs[0]->params["__batch_index"].i;
3737

3838
int axis = captured_params.at("dim").i;
39-
if (axis == batch_index)
39+
if (axis < 0)
4040
{
41-
fprintf(stderr, "glu along batch axis %d is not supported\n", batch_index);
42-
return;
41+
int input_rank = op->inputs[0]->shape.size();
42+
if (input_rank > 0)
43+
axis = input_rank + axis;
4344
}
4445

45-
if (axis < 0)
46+
if (axis == batch_index)
4647
{
47-
int input_rank = op->inputs[0]->shape.size();
48-
axis = input_rank + axis;
48+
fprintf(stderr, "glu along batch axis %d is not supported\n", batch_index);
49+
return;
4950
}
5051

5152
if (axis > batch_index)

‎tools/pnnx/src/pass_ncnn/F_grid_sample.cpp‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,11 +116,13 @@ pnnx.Output output 1 0 out
116116
const std::vector<int>& dims = captured_params.at("dims").ai;
117117

118118
int input_rank = (int)op->inputs[0]->shape.size();
119+
int full_input_rank = input_rank;
119120

120121
if (input_rank == 0)
121122
{
122123
// assume input is fine
123124
input_rank = (int)dims.size();
125+
full_input_rank = input_rank;
124126
}
125127

126128
if (batch_index >= 0 && batch_index < input_rank)
@@ -136,10 +138,14 @@ pnnx.Output output 1 0 out
136138
std::vector<int> new_dims;
137139
for (int i = 0; i < (int)dims.size(); i++)
138140
{
139-
if (dims[i] == batch_index)
141+
int dim = dims[i];
142+
if (dim < 0 && full_input_rank > 0)
143+
dim += full_input_rank;
144+
145+
if (dim == batch_index)
140146
continue;
141147

142-
int new_dim = dims[i] > batch_index ? dims[i] - 1 : dims[i];
148+
int new_dim = dim > batch_index ? dim - 1 : dim;
143149
new_dims.push_back(new_dim);
144150
}
145151

‎tools/pnnx/src/pass_ncnn/F_normalize.cpp‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,16 +35,17 @@ pnnx.Output output 1 0 out
3535
const int batch_index = op->inputs[0]->params["__batch_index"].i;
3636

3737
int axis = captured_params.at("dim").i;
38-
if (axis == batch_index)
38+
if (axis < 0)
3939
{
40-
fprintf(stderr, "normalize along batch axis %d is not supported\n", batch_index);
41-
return;
40+
int input_rank = op->inputs[0]->shape.size();
41+
if (input_rank > 0)
42+
axis = input_rank + axis;
4243
}
4344

44-
if (axis < 0)
45+
if (axis == batch_index)
4546
{
46-
int input_rank = op->inputs[0]->shape.size();
47-
axis = input_rank + axis;
47+
fprintf(stderr, "normalize along batch axis %d is not supported\n", batch_index);
48+
return;
4849
}
4950

5051
if (axis > batch_index)

‎tools/pnnx/src/pass_ncnn/F_softmax.cpp‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,16 +35,17 @@ pnnx.Output output 1 0 out
3535
const int batch_index = op->inputs[0]->params["__batch_index"].i;
3636

3737
int axis = captured_params.at("dim").i;
38-
if (axis == batch_index)
38+
if (axis < 0)
3939
{
40-
fprintf(stderr, "softmax along batch axis %d is not supported\n", batch_index);
41-
return;
40+
int input_rank = op->inputs[0]->shape.size();
41+
if (input_rank > 0)
42+
axis = input_rank + axis;
4243
}
4344

44-
if (axis < 0)
45+
if (axis == batch_index)
4546
{
46-
int input_rank = op->inputs[0]->shape.size();
47-
axis = input_rank + axis;
47+
fprintf(stderr, "softmax along batch axis %d is not supported\n", batch_index);
48+
return;
4849
}
4950

5051
if (axis > batch_index)

‎tools/pnnx/src/pass_ncnn/Tensor_permute.cpp‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,11 +39,13 @@ pnnx.Output output 1 0 out
3939
const std::vector<int>& dims = captured_params.at("dims").ai;
4040

4141
int input_rank = (int)op->inputs[0]->shape.size();
42+
int full_input_rank = input_rank;
4243

4344
if (input_rank == 0)
4445
{
4546
// assume input is fine
4647
input_rank = (int)dims.size();
48+
full_input_rank = input_rank;
4749
}
4850

4951
if (batch_index >= 0 && batch_index < input_rank)
@@ -59,10 +61,14 @@ pnnx.Output output 1 0 out
5961
std::vector<int> new_dims;
6062
for (int i = 0; i < (int)dims.size(); i++)
6163
{
62-
if (dims[i] == batch_index)
64+
int dim = dims[i];
65+
if (dim < 0 && full_input_rank > 0)
66+
dim += full_input_rank;
67+
68+
if (dim == batch_index)
6369
continue;
6470

65-
int new_dim = dims[i] > batch_index ? dims[i] - 1 : dims[i];
71+
int new_dim = dim > batch_index ? dim - 1 : dim;
6672
new_dims.push_back(new_dim);
6773
}
6874

‎tools/pnnx/src/pass_ncnn/convert_Tensor_select.cpp‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,16 +28,17 @@ void convert_Tensor_select(Graph& graph)
2828
const int batch_index = op->inputs[0]->params["__batch_index"].i;
2929

3030
int axis = op->params.at("dim").i;
31-
if (axis == batch_index)
31+
if (axis < 0)
3232
{
33-
fprintf(stderr, "select along batch axis %d is not supported\n", batch_index);
34-
continue;
33+
int input_rank = op->inputs[0]->shape.size();
34+
if (input_rank > 0)
35+
axis = input_rank + axis;
3536
}
3637

37-
if (axis < 0)
38+
if (axis == batch_index)
3839
{
39-
int input_rank = op->inputs[0]->shape.size();
40-
axis = input_rank + axis;
40+
fprintf(stderr, "select along batch axis %d is not supported\n", batch_index);
41+
continue;
4142
}
4243

4344
if (axis > batch_index)
@@ -47,7 +48,9 @@ void convert_Tensor_select(Graph& graph)
4748
int index;
4849
if (op->has_param("dim"))
4950
{
50-
dim = op->params.at("dim").i;
51+
dim = axis;
52+
if (dim >= batch_index)
53+
dim += 1;
5154
}
5255
else
5356
{

‎tools/pnnx/src/pass_ncnn/convert_Tensor_slice.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -140,18 +140,18 @@ void convert_Tensor_slice(Graph& graph)
140140
}
141141
}
142142

143+
int input_rank0 = op->inputs[0]->shape.size();
143144
for (int i = 0; i < axes_rank; i++)
144145
{
145-
if (axes[i] == batch_index && (starts[i] != 0 || ends[i] != INT_MAX))
146+
if (axes[i] < 0 && input_rank0 > 0)
146147
{
147-
fprintf(stderr, "slice along batch axis is not supported\n");
148-
continue;
148+
axes[i] = input_rank0 + axes[i];
149149
}
150150

151-
if (axes[i] < 0)
151+
if (axes[i] == batch_index && (starts[i] != 0 || ends[i] != INT_MAX))
152152
{
153-
int input_rank = op->inputs[0]->shape.size();
154-
axes[i] = input_rank + axes[i];
153+
fprintf(stderr, "slice along batch axis is not supported\n");
154+
continue;
155155
}
156156

157157
if (axes[i] > batch_index)

‎tools/pnnx/src/pass_ncnn/convert_Tensor_slice_copy.cpp‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ void convert_Tensor_slice_copy(Graph& graph)
108108
const int axes_rank = axes.size();
109109

110110
bool has_select = false;
111-
std::vector<int> selected_axes;
111+
std::vector<int> selected_axis_indices;
112112
for (int i = 0; i < axes_rank; i++)
113113
{
114114
if (steps[i] == 0)
@@ -118,7 +118,7 @@ void convert_Tensor_slice_copy(Graph& graph)
118118
ends[i] = selects[i] + 1;
119119
steps[i] = 1;
120120
has_select = true;
121-
selected_axes.push_back(axes[i]);
121+
selected_axis_indices.push_back(i);
122122
}
123123
else if (steps[i] != 1)
124124
{
@@ -141,18 +141,18 @@ void convert_Tensor_slice_copy(Graph& graph)
141141
}
142142
}
143143

144+
int input_rank0 = op->inputs[0]->shape.size();
144145
for (int i = 0; i < axes_rank; i++)
145146
{
146-
if (axes[i] == batch_index && (starts[i] != 0 || ends[i] != INT_MAX))
147+
if (axes[i] < 0 && input_rank0 > 0)
147148
{
148-
fprintf(stderr, "slice_copy along batch axis is not supported\n");
149-
continue;
149+
axes[i] = input_rank0 + axes[i];
150150
}
151151

152-
if (axes[i] < 0)
152+
if (axes[i] == batch_index && (starts[i] != 0 || ends[i] != INT_MAX))
153153
{
154-
int input_rank = op->inputs[0]->shape.size();
155-
axes[i] = input_rank + axes[i];
154+
fprintf(stderr, "slice_copy along batch axis is not supported\n");
155+
continue;
156156
}
157157

158158
if (axes[i] > batch_index)
@@ -203,9 +203,15 @@ void convert_Tensor_slice_copy(Graph& graph)
203203
in->consumers.push_back(reshape);
204204

205205
std::vector<int> shape = in->shape;
206-
for (auto sa : selected_axes)
206+
for (auto si : selected_axis_indices)
207207
{
208208
// unsqueeze
209+
int sa = axes[si];
210+
if (shape.empty())
211+
continue;
212+
if (sa < 0 || sa > (int)shape.size())
213+
continue;
214+
209215
shape.insert(shape.begin() + sa, 1);
210216
}
211217

‎tools/pnnx/src/pass_ncnn/convert_torch_cat.cpp‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,17 @@ void convert_torch_cat(Graph& graph)
2222
const int batch_index = op->inputs[0]->params["__batch_index"].i;
2323

2424
int axis = op->params.at("dim").i;
25-
if (axis == batch_index)
25+
if (axis < 0)
2626
{
27-
fprintf(stderr, "cat along batch axis %d is not supported\n", batch_index);
28-
continue;
27+
int input_rank = op->inputs[0]->shape.size();
28+
if (input_rank > 0)
29+
axis = input_rank + axis;
2930
}
3031

31-
if (axis < 0)
32+
if (axis == batch_index)
3233
{
33-
int input_rank = op->inputs[0]->shape.size();
34-
axis = input_rank + axis;
34+
fprintf(stderr, "cat along batch axis %d is not supported\n", batch_index);
35+
continue;
3536
}
3637

3738
if (axis > batch_index)

‎tools/pnnx/src/pass_ncnn/convert_torch_chunk.cpp‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,17 @@ void convert_torch_chunk(Graph& graph)
2222
const int batch_index = op->inputs[0]->params["__batch_index"].i;
2323

2424
int axis = op->params.at("dim").i;
25-
if (axis == batch_index)
25+
if (axis < 0)
2626
{
27-
fprintf(stderr, "chunk along batch axis %d is not supported\n", batch_index);
28-
continue;
27+
int input_rank = op->inputs[0]->shape.size();
28+
if (input_rank > 0)
29+
axis = input_rank + axis;
2930
}
3031

31-
if (axis < 0)
32+
if (axis == batch_index)
3233
{
33-
int input_rank = op->inputs[0]->shape.size();
34-
axis = input_rank + axis;
34+
fprintf(stderr, "chunk along batch axis %d is not supported\n", batch_index);
35+
continue;
3536
}
3637

3738
int chunks = op->params.at("chunks").i;

0 commit comments

Comments
 (0)