lib/torch/nn/functional.rb in torch-rb-0.3.6 vs lib/torch/nn/functional.rb in torch-rb-0.3.7

@@ -467,10 +467,81 @@
 
         def triplet_margin_loss(anchor, positive, negative, margin: 1.0, p: 2, eps: 1e-06, swap: false, reduction: "mean")
           Torch.triplet_margin_loss(anchor, positive, negative, margin, p, eps, swap, reduction)
         end
 
+        # vision
+
+        def interpolate(input, size: nil, scale_factor: nil, mode: "nearest", align_corners: nil, recompute_scale_factor: nil)
+          if ["nearest", "area"].include?(mode)
+            unless align_corners.nil?
+              raise ArgumentError, "align_corners option can only be set with the interpolating modes: linear | bilinear | bicubic | trilinear"
+            end
+          else
+            if align_corners.nil?
+              align_corners = false
+            end
+          end
+
+          scale_factor_len = input.dim - 2
+          scale_factor_list = [nil] * scale_factor_len
+          # default value of recompute_scale_factor is False
+          if !scale_factor.nil? && (recompute_scale_factor == false || recompute_scale_factor.nil?)
+            if scale_factor.is_a?(Array)
+              _scale_factor_repeated = scale_factor
+            else
+              _scale_factor_repeated = [scale_factor] * scale_factor_len
+            end
+            scale_factor_list = _scale_factor_repeated
+          end
+
+          # Give this variable a short name because it has to be repeated multiple times below.
+          sfl = scale_factor_list
+
+          closed_over_args = [input, size, scale_factor, recompute_scale_factor]
+          output_size = _interp_output_size(closed_over_args)
+          if input.dim == 3 && mode == "nearest"
+            NN.upsample_nearest1d(input, output_size, sfl[0])
+          elsif input.dim == 4 && mode == "nearest"
+            NN.upsample_nearest2d(input, output_size, sfl[0], sfl[1])
+          elsif input.dim == 5 && mode == "nearest"
+            NN.upsample_nearest3d(input, output_size, sfl[0], sfl[1], sfl[2])
+          elsif input.dim == 3 && mode == "area"
+            adaptive_avg_pool1d(input, output_size)
+          elsif input.dim == 4 && mode == "area"
+            adaptive_avg_pool2d(input, output_size)
+          elsif input.dim == 5 && mode == "area"
+            adaptive_avg_pool3d(input, output_size)
+          elsif input.dim == 3 && mode == "linear"
+            # assert align_corners is not None
+            NN.upsample_linear1d(input, output_size, align_corners, sfl[0])
+          elsif input.dim == 3 && mode == "bilinear"
+            raise ArgumentError, "Got 3D input, but bilinear mode needs 4D input"
+          elsif input.dim == 3 && mode == "trilinear"
+            raise ArgumentError, "Got 3D input, but trilinear mode needs 5D input"
+          elsif input.dim == 4 && mode == "linear"
+            raise ArgumentError, "Got 4D input, but linear mode needs 3D input"
+          elsif input.dim == 4 && mode == "bilinear"
+            # assert align_corners is not None
+            NN.upsample_bilinear2d(input, output_size, align_corners, sfl[0], sfl[1])
+          elsif input.dim == 4 && mode == "trilinear"
+            raise ArgumentError, "Got 4D input, but trilinear mode needs 5D input"
+          elsif input.dim == 5 && mode == "linear"
+            raise ArgumentError, "Got 5D input, but linear mode needs 3D input"
+          elsif input.dim == 5 && mode == "bilinear"
+            raise ArgumentError, "Got 5D input, but bilinear mode needs 4D input"
+          elsif input.dim == 5 && mode == "trilinear"
+            # assert align_corners is not None
+            NN.upsample_trilinear3d(input, output_size, align_corners, sfl[0], sfl[1], sfl[2])
+          elsif input.dim == 4 && mode == "bicubic"
+            # assert align_corners is not None
+            NN.upsample_bicubic2d(input, output_size, align_corners, sfl[0], sfl[1])
+          else
+            raise ArgumentError, "Input Error: Only 3D, 4D and 5D input Tensors supported (got #{input.dim}D) for the modes: nearest | linear | bilinear | bicubic | trilinear (got #{mode})"
+          end
+        end
+
         private
 
         def softmax_dim(ndim)
           ndim == 0 || ndim == 1 || ndim == 3 ? 0 : 1
         end
@@ -481,9 +552,44 @@
           elsif defaults.length < out_size.length
             raise ArgumentError, "Input dimension should be at least #{out_size.length + 1}"
           else
             out_size.zip(defaults.last(out_size.length)).map { |v, d| v || d }
           end
+        end
+
+        def _interp_output_size(closed_over_args)
+          input, size, scale_factor, recompute_scale_factor = closed_over_args
+          dim = input.dim - 2
+          if size.nil? && scale_factor.nil?
+            raise ArgumentError, "either size or scale_factor should be defined"
+          end
+          if !size.nil? && !scale_factor.nil?
+            raise ArgumentError, "only one of size or scale_factor should be defined"
+          end
+          if !scale_factor.nil?
+            if scale_factor.is_a?(Array)
+              if scale_factor.length != dim
+                raise ArgumentError, "scale_factor shape must match input shape. Input is #{dim}D, scale_factor size is #{scale_factor.length}"
+              end
+            end
+          end
+
+          if !size.nil?
+            if size.is_a?(Array)
+              return size
+            else
+              return [size] * dim
+            end
+          end
+
+          raise "Failed assertion" if scale_factor.nil?
+          if scale_factor.is_a?(Array)
+            scale_factors = scale_factor
+          else
+            scale_factors = [scale_factor] * dim
+          end
+
+          dim.times.map { |i| (input.size(i + 2) * scale_factors[i]).floor }
         end
       end
     end
 
     # shortcut