torch-rb 0.8.1 → 0.9.1

data/lib/torch/nn/parameter.rb

@@ -9,6 +9,12 @@ module Torch
       def inspect
         "Parameter containing:\n#{super}"
       end
+
+      def dup
+        Torch.no_grad do
+          Parameter.new(clone, requires_grad: requires_grad)
+        end
+      end
     end
   end
 end

data/lib/torch/nn/transformer.rb

@@ -0,0 +1,92 @@
+require_relative 'transformer_encoder_layer'
+require_relative 'transformer_encoder'
+require_relative 'transformer_decoder_layer'
+require_relative 'transformer_decoder'
+
+module Torch
+  module NN
+    class Transformer < Module
+      def initialize(
+        d_model: 512, nhead: 8,
+        num_encoder_layers: 6, num_decoder_layers: 6,
+        dim_feedforward: 2048, dropout: 0.1, activation: :relu,
+        custom_encoder: nil, custom_decoder: nil,
+        layer_norm_eps: 1e-5, batch_first: false
+      )
+
+        super()
+
+        @encoder =
+          if custom_encoder
+            custom_encoder
+          else
+            encoder_layer = TransformerEncoderLayer.new(
+              d_model, nhead,
+              dim_feedforward: dim_feedforward, dropout: dropout, activation: activation,
+              layer_norm_eps: layer_norm_eps, batch_first: batch_first
+            )
+            encoder_norm = LayerNorm.new(d_model, eps: layer_norm_eps)
+            TransformerEncoder.new(encoder_layer, num_encoder_layers, norm: encoder_norm)
+          end
+
+        @decoder =
+          if custom_decoder
+            custom_decoder
+          else
+            decoder_layer = TransformerDecoderLayer.new(
+              d_model, nhead,
+              dim_feedforward: dim_feedforward, dropout: dropout, activation: activation,
+              layer_norm_eps: layer_norm_eps, batch_first: batch_first
+            )
+            decoder_norm = LayerNorm.new(d_model, eps: layer_norm_eps)
+            TransformerDecoder.new(decoder_layer, num_decoder_layers, norm: decoder_norm)
+          end
+
+        reset_parameters
+
+        @d_model = d_model
+        @nhead = nhead
+        @batch_first = batch_first
+      end
+
+      attr_reader :d_model, :nhead, :encoder, :decoder
+
+      def batch_first?
+        !!@batch_first
+      end
+
+      def reset_parameters
+        parameters.each { |p| Init.xavier_uniform!(p) if p.dim > 1 }
+      end
+
+      def forward(
+        src, tgt,
+        src_mask: nil, tgt_mask: nil, memory_mask: nil,
+        src_key_padding_mask: nil, tgt_key_padding_mask: nil, memory_key_padding_mask: nil
+      )
+
+        if (!batch_first? && src.size(1) != tgt.size(1)) ||
+          (batch_first? && src.size(0) != tgt.size(0))
+
+          raise ArgumentError, "The batch number of src and tgt must be equal"
+        end
+
+        if src.size(2) != d_model || tgt.size(2) != d_model
+          raise ArgumentError, "The feature number of src and tgt must be equal to d_model"
+        end
+
+        memory = @encoder.(src, mask: src_mask, src_key_padding_mask: src_key_padding_mask)
+        @decoder.(
+          tgt, memory,
+          tgt_mask: tgt_mask, memory_mask: memory_mask,
+          tgt_key_padding_mask: tgt_key_padding_mask, memory_key_padding_mask: memory_key_padding_mask
+        )
+      end
+
+      def generate_square_subsequent_mask(sz)
+        mask = Torch.triu(Torch.ones([sz, sz])).eq(1).transpose(0, 1)
+        mask.float.masked_fill!(mask.eq(0), -Float::INFINITY).masked_fill!(mask.eq(1), 0.0)
+      end
+    end
+  end
+end

data/lib/torch/nn/transformer_decoder.rb

@@ -0,0 +1,25 @@
+module Torch
+  module NN
+    class TransformerDecoder < Module
+      def initialize(decoder_layer, num_layers, norm: nil)
+        super()
+
+        @layers = _clones(decoder_layer, num_layers)
+        @num_layers = num_layers
+        @norm = norm
+      end
+
+      def forward(tgt, memory, tgt_mask: nil, memory_mask: nil, tgt_key_padding_mask: nil, memory_key_padding_mask: nil)
+        output = tgt
+
+        @layers.each do |mod|
+          output = mod.call(output, memory, tgt_mask: tgt_mask, memory_mask: memory_mask, tgt_key_padding_mask: tgt_key_padding_mask, memory_key_padding_mask: memory_key_padding_mask)
+        end
+
+        output = @norm.call(output) if @norm
+
+        output
+      end
+    end
+  end
+end

data/lib/torch/nn/transformer_decoder_layer.rb

@@ -0,0 +1,43 @@
+module Torch
+  module NN
+    class TransformerDecoderLayer < Module
+      def initialize(
+        d_model, n_head,
+        dim_feedforward: 2048, dropout: 0.1, activation: :relu,
+        layer_norm_eps: 1e-5, batch_first: false
+      )
+
+        super()
+
+        @self_attn = MultiheadAttention.new(d_model, n_head, dropout: dropout, batch_first: batch_first)
+        @multihead_attn = MultiheadAttention.new(d_model, n_head, dropout: dropout, batch_first: batch_first)
+
+        @linear1 = Linear.new(d_model, dim_feedforward)
+        @dropout = Dropout.new(p: dropout)
+        @linear2 = Linear.new(dim_feedforward, d_model)
+
+        @norm1 = LayerNorm.new(d_model, eps: layer_norm_eps)
+        @norm2 = LayerNorm.new(d_model, eps: layer_norm_eps)
+        @norm3 = LayerNorm.new(d_model, eps: layer_norm_eps)
+
+        @dropout1 = Dropout.new(p: dropout)
+        @dropout2 = Dropout.new(p: dropout)
+        @dropout3 = Dropout.new(p: dropout)
+
+        @activation = _activation_fn(activation)
+      end
+
+      def forward(tgt, memory, tgt_mask: nil, memory_mask: nil, tgt_key_padding_mask: nil, memory_key_padding_mask: nil)
+        tgt2 = @self_attn.(tgt, tgt, tgt, attn_mask: tgt_mask, key_padding_mask: tgt_key_padding_mask).first
+        tgt += @dropout1.(tgt2)
+        tgt = @norm1.(tgt)
+        tgt2 = @multihead_attn.(tgt, memory, memory, attn_mask: memory_mask, key_padding_mask: memory_key_padding_mask).first
+        tgt += @dropout2.(tgt2)
+        tgt = @norm2.(tgt)
+        tgt2 = @linear2.(@dropout.(@activation.(@linear1.(tgt))))
+        tgt += @dropout3.(tgt2)
+        @norm3.(tgt)
+      end
+    end
+  end
+end

data/lib/torch/nn/transformer_encoder.rb

@@ -0,0 +1,25 @@
+module Torch
+  module NN
+    class TransformerEncoder < Module
+      def initialize(encoder_layer, num_layers, norm: nil)
+        super()
+
+        @layers = _clones(encoder_layer, num_layers)
+        @num_layers = num_layers
+        @norm = norm
+      end
+
+      def forward(src, mask: nil, src_key_padding_mask: nil)
+        output = src
+
+        @layers.each do |mod|
+          output = mod.call(output, src_mask: mask, src_key_padding_mask: src_key_padding_mask)
+        end
+
+        output = @norm.call(output) if @norm
+
+        output
+      end
+    end
+  end
+end

data/lib/torch/nn/transformer_encoder_layer.rb

@@ -0,0 +1,36 @@
+module Torch
+  module NN
+    class TransformerEncoderLayer < Module
+      def initialize(
+        d_model, n_head,
+        dim_feedforward: 2048, dropout: 0.1, activation: :relu,
+        layer_norm_eps: 1e-5, batch_first: false
+      )
+
+        super()
+
+        @self_attn = MultiheadAttention.new(d_model, n_head, dropout: dropout, batch_first: batch_first)
+        @linear1 = Linear.new(d_model, dim_feedforward)
+        @dropout = Dropout.new(p: dropout)
+        @linear2 = Linear.new(dim_feedforward, d_model)
+
+        @norm1 = LayerNorm.new(d_model, eps: layer_norm_eps)
+        @norm2 = LayerNorm.new(d_model, eps: layer_norm_eps)
+
+        @dropout1 = Dropout.new(p: dropout)
+        @dropout2 = Dropout.new(p: dropout)
+
+        @activation = _activation_fn(activation)
+      end
+
+      def forward(src, src_mask: nil, src_key_padding_mask: nil)
+        src2 = @self_attn.(src, src, src, attn_mask: src_mask, key_padding_mask: src_key_padding_mask).first
+        src += @dropout1.(src2)
+        src = @norm1.(src)
+        src2 = @linear2.(@dropout.(@activation.(@linear1.(src))))
+        src += @dropout2.(src2)
+        @norm2.(src)
+      end
+    end
+  end
+end

data/lib/torch/nn/utils.rb

@@ -20,6 +20,18 @@ module Torch
       def _ntuple(n, value)
         value.is_a?(Array) ? value : [value] * n
       end
+
+      def _clones(mod, n)
+        ModuleList.new(n.times.map { mod.deep_dup })
+      end
+
+      def _activation_fn(activation)
+        case activation.to_sym
+        when :relu then F.method(:relu)
+        when :gelu then F.method(:gelu)
+        else raise ArgumentError, "Activation should be relu/gelu, not `#{activation}`"
+        end
+      end
     end
   end
 end

data/lib/torch/tensor.rb

@@ -19,6 +19,8 @@ module Torch
     alias_method :&, :logical_and
     alias_method :|, :logical_or
     alias_method :^, :logical_xor
+    alias_method :<<, :__lshift__
+    alias_method :>>, :__rshift__
 
     def self.new(*args)
       FloatTensor.new(*args)
@@ -104,6 +106,7 @@ module Torch
       size(0)
     end
 
+    remove_method :item
     def item
       if numel != 1
         raise Error, "only one element tensors can be converted to Ruby scalars"
@@ -131,18 +134,10 @@ module Torch
       cls.from_string(_data_str).reshape(*shape)
     end
 
-    def new_ones(*size, **options)
-      Torch.ones_like(Torch.empty(*size), **options)
-    end
-
     def requires_grad=(requires_grad)
       _requires_grad!(requires_grad)
     end
 
-    def requires_grad!(requires_grad = true)
-      _requires_grad!(requires_grad)
-    end
-
     def type(dtype)
       if dtype.is_a?(Class)
         raise Error, "Invalid type: #{dtype}" unless TENSOR_TYPE_CLASSES.include?(dtype)
@@ -183,5 +178,23 @@ module Torch
     def stft(*args)
       Torch.stft(*args)
     end
+
+    def dup
+      Torch.no_grad do
+        clone
+      end
+    end
+
+    # not a method in native_functions.yaml
+    # attribute in Python rather than method
+    def imag
+      Torch.imag(self)
+    end
+
+    # not a method in native_functions.yaml
+    # attribute in Python rather than method
+    def real
+      Torch.real(self)
+    end
   end
 end

data/lib/torch/utils/data/data_loader.rb

@@ -25,9 +25,11 @@ module Torch
         end
 
         def each
+          return to_enum(:each) unless block_given?
+
           # try to keep the random number generator in sync with Python
           # this makes it easy to compare results
-          base_seed = Torch.empty([], dtype: :int64).random!.item
+          _base_seed = Torch.empty([], dtype: :int64).random!.item
 
           indexes =
             if @shuffle

data/lib/torch/version.rb

@@ -1,3 +1,3 @@
 module Torch
-  VERSION = "0.8.1"
+  VERSION = "0.9.1"
 end

data/lib/torch.rb

@@ -39,6 +39,7 @@ require "torch/nn/utils"
 
 # nn containers
 require "torch/nn/module"
+require "torch/nn/module_list"
 require "torch/nn/sequential"
 
 # nn convolution layers
@@ -143,6 +144,10 @@ require "torch/nn/softmin"
 require "torch/nn/embedding"
 require "torch/nn/embedding_bag"
 
+# attention is all you need
+require "torch/nn/multihead_attention"
+require "torch/nn/transformer"
+
 # nn distance functions
 require "torch/nn/cosine_similarity"
 require "torch/nn/pairwise_distance"
@@ -174,6 +179,7 @@ require "torch/nn/upsample"
 
 # nn other
 require "torch/nn/functional"
+require "torch/nn/functional_attention"
 require "torch/nn/init"
 
 # utils
@@ -371,8 +377,6 @@ module Torch
       to_ruby(_load(File.binread(f)))
     end
 
-    # --- begin tensor creation: https://pytorch.org/cppdocs/notes/tensor_creation.html ---
-
     def tensor(data, **options)
       if options[:dtype].nil? && defined?(Numo::NArray) && data.is_a?(Numo::NArray)
         numo_to_dtype = _dtype_to_numo.map(&:reverse).to_h
@@ -405,41 +409,6 @@ module Torch
       _tensor(data, size, tensor_options(**options))
     end
 
-    # --- begin like ---
-
-    def ones_like(input, **options)
-      ones(input.size, **like_options(input, options))
-    end
-
-    def empty_like(input, **options)
-      empty(input.size, **like_options(input, options))
-    end
-
-    def full_like(input, fill_value, **options)
-      full(input.size, fill_value, **like_options(input, options))
-    end
-
-    def rand_like(input, **options)
-      rand(input.size, **like_options(input, options))
-    end
-
-    def randint_like(input, low, high = nil, **options)
-      # ruby doesn't support input, low = 0, high, ...
-      if high.nil?
-        high = low
-        low = 0
-      end
-      randint(low, high, input.size, **like_options(input, options))
-    end
-
-    def randn_like(input, **options)
-      randn(input.size, **like_options(input, options))
-    end
-
-    def zeros_like(input, **options)
-      zeros(input.size, **like_options(input, options))
-    end
-
     # center option
     def stft(input, n_fft, hop_length: nil, win_length: nil, window: nil, center: true, pad_mode: "reflect", normalized: false, onesided: true, return_complex: nil)
       if center
@@ -566,13 +535,5 @@ module Torch
       end
       options
     end
-
-    def like_options(input, options)
-      options = options.dup
-      options[:dtype] ||= input.dtype
-      options[:layout] ||= input.layout
-      options[:device] ||= input.device
-      options
-    end
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: torch-rb
 version: !ruby/object:Gem::Version
-  version: 0.8.1
+  version: 0.9.1
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-06-16 00:00:00.000000000 Z
+date: 2022-02-03 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice
@@ -106,6 +106,7 @@ files:
 - lib/torch/nn/feature_alpha_dropout.rb
 - lib/torch/nn/fold.rb
 - lib/torch/nn/functional.rb
+- lib/torch/nn/functional_attention.rb
 - lib/torch/nn/group_norm.rb
 - lib/torch/nn/gru.rb
 - lib/torch/nn/hardshrink.rb
@@ -139,10 +140,12 @@ files:
 - lib/torch/nn/max_unpool3d.rb
 - lib/torch/nn/max_unpoolnd.rb
 - lib/torch/nn/module.rb
+- lib/torch/nn/module_list.rb
 - lib/torch/nn/mse_loss.rb
 - lib/torch/nn/multi_label_margin_loss.rb
 - lib/torch/nn/multi_label_soft_margin_loss.rb
 - lib/torch/nn/multi_margin_loss.rb
+- lib/torch/nn/multihead_attention.rb
 - lib/torch/nn/nll_loss.rb
 - lib/torch/nn/pairwise_distance.rb
 - lib/torch/nn/parameter.rb
@@ -170,6 +173,11 @@ files:
 - lib/torch/nn/softsign.rb
 - lib/torch/nn/tanh.rb
 - lib/torch/nn/tanhshrink.rb
+- lib/torch/nn/transformer.rb
+- lib/torch/nn/transformer_decoder.rb
+- lib/torch/nn/transformer_decoder_layer.rb
+- lib/torch/nn/transformer_encoder.rb
+- lib/torch/nn/transformer_encoder_layer.rb
 - lib/torch/nn/triplet_margin_loss.rb
 - lib/torch/nn/unfold.rb
 - lib/torch/nn/upsample.rb
@@ -219,7 +227,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.2.3
+rubygems_version: 3.3.3
 signing_key:
 specification_version: 4
 summary: Deep learning for Ruby, powered by LibTorch