stu1130 commented on a change in pull request #13614: Make to_tensor and
normalize to accept 3D or 4D tensor inputs
URL: https://github.com/apache/incubator-mxnet/pull/13614#discussion_r241201309
##########
File path: src/operator/image/image_random-inl.h
##########
@@ -70,19 +92,23 @@ void ToTensor(const nnvm::NodeAttrs &attrs,
CHECK_EQ(req[0], kWriteTo)
<< "`to_tensor` does not support inplace";
- int length = inputs[0].shape_[0] * inputs[0].shape_[1];
- int channel = inputs[0].shape_[2];
-
- MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, DType, {
- float* output = outputs[0].dptr<float>();
- DType* input = inputs[0].dptr<DType>();
-
- for (int l = 0; l < length; ++l) {
- for (int c = 0; c < channel; ++c) {
- output[c*length + l] = static_cast<float>(input[l*channel + c]) /
255.0f;
- }
+ // 3D Input - 1 image
+ if (inputs[0].ndim() == 3) {
+ const int length = inputs[0].shape_[0] * inputs[0].shape_[1];
+ const int channel = inputs[0].shape_[2];
+ ToTensorImpl(inputs, outputs, length, channel);
+ } else if (inputs[0].ndim() == 4) {
+ // 4D input batch of images
+ const int batch_size = inputs[0].shape_[0];
+ const int length = inputs[0].shape_[1] * inputs[0].shape_[2];
+ const int channel = inputs[0].shape_[3];
+ const int step = channel * length;
Review comment:
Good point, IMO make step unsigned long long int would be enough?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
