torchvision 0.1.1 → 0.2.2

data/lib/torchvision/datasets/dataset_folder.rb

@@ -0,0 +1,91 @@
+module TorchVision
+  module Datasets
+    class DatasetFolder < VisionDataset
+      attr_reader :classes
+
+      def initialize(root, extensions: nil, transform: nil, target_transform: nil, is_valid_file: nil)
+        super(root, transform: transform, target_transform: target_transform)
+        classes, class_to_idx = find_classes(@root)
+        samples = make_dataset(@root, class_to_idx, extensions, is_valid_file)
+        if samples.empty?
+          msg = "Found 0 files in subfolders of: #{@root}\n"
+          unless extensions.nil?
+            msg += "Supported extensions are: #{extensions.join(",")}"
+          end
+          raise RuntimeError, msg
+        end
+
+        @loader = lambda do |path|
+          Vips::Image.new_from_file(path)
+        end
+        @extensions = extensions
+
+        @classes = classes
+        @class_to_idx = class_to_idx
+        @samples = samples
+        @targets = samples.map { |s| s[1] }
+      end
+
+      def [](index)
+        path, target = @samples[index]
+        sample = @loader.call(path)
+        if @transform
+          sample = @transform.call(sample)
+        end
+        if @target_transform
+          target = @target_transform.call(target)
+        end
+
+        [sample, target]
+      end
+
+      def size
+        @samples.size
+      end
+
+      private
+
+      def find_classes(dir)
+        classes = Dir.children(dir).select { |d| File.directory?(File.join(dir, d)) }
+        classes.sort!
+        class_to_idx = classes.map.with_index.to_h
+        [classes, class_to_idx]
+      end
+
+      def has_file_allowed_extension(filename, extensions)
+        filename = filename.downcase
+        extensions.any? { |ext| filename.end_with?(ext) }
+      end
+
+      def make_dataset(directory, class_to_idx, extensions, is_valid_file)
+        instances = []
+        directory = File.expand_path(directory)
+        both_none = extensions.nil? && is_valid_file.nil?
+        both_something = !extensions.nil? && !is_valid_file.nil?
+        if both_none || both_something
+          raise ArgumentError, "Both extensions and is_valid_file cannot be None or not None at the same time"
+        end
+        if !extensions.nil?
+          is_valid_file = lambda do |x|
+            has_file_allowed_extension(x, extensions)
+          end
+        end
+        class_to_idx.keys.sort.each do |target_class|
+          class_index = class_to_idx[target_class]
+          target_dir = File.join(directory, target_class)
+          if !File.directory?(target_dir)
+            next
+          end
+          Dir.glob("**", base: target_dir).sort.each do |fname|
+            path = File.join(target_dir, fname)
+            if is_valid_file.call(path)
+              item = [path, class_index]
+              instances << item
+            end
+          end
+        end
+        instances
+      end
+    end
+  end
+end

data/lib/torchvision/datasets/fashion_mnist.rb

@@ -0,0 +1,30 @@
+module TorchVision
+  module Datasets
+    class FashionMNIST < MNIST
+      # https://github.com/zalandoresearch/fashion-mnist
+
+      private
+
+      def resources
+        [
+          {
+            url: "http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz",
+            sha256: "3aede38d61863908ad78613f6a32ed271626dd12800ba2636569512369268a84"
+          },
+          {
+            url: "http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz",
+            sha256: "a04f17134ac03560a47e3764e11b92fc97de4d1bfaf8ba1a3aa29af54cc90845"
+          },
+          {
+            url: "http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz",
+            sha256: "346e55b948d973a97e58d2351dde16a484bd415d4595297633bb08f03db6a073"
+          },
+          {
+            url: "http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz",
+            sha256: "67da17c76eaffca5446c3361aaab5c3cd6d1c2608764d35dfb1850b086bf8dd5"
+          }
+        ]
+      end
+    end
+  end
+end

data/lib/torchvision/datasets/image_folder.rb

@@ -0,0 +1,12 @@
+module TorchVision
+  module Datasets
+    class ImageFolder < DatasetFolder
+      IMG_EXTENSIONS = [".jpg", ".jpeg", ".png", ".ppm", ".bmp", ".pgm", ".tif", ".tiff", ".webp"]
+
+      def initialize(root, transform: nil, target_transform: nil, is_valid_file: nil)
+        super(root, extensions: IMG_EXTENSIONS, transform: transform, target_transform: target_transform, is_valid_file: is_valid_file)
+        @imgs = @samples
+      end
+    end
+  end
+end

data/lib/torchvision/datasets/kmnist.rb

@@ -0,0 +1,30 @@
+module TorchVision
+  module Datasets
+    class KMNIST < MNIST
+      # https://github.com/rois-codh/kmnist
+
+      private
+
+      def resources
+        [
+          {
+            url: "http://codh.rois.ac.jp/kmnist/dataset/kmnist/train-images-idx3-ubyte.gz",
+            sha256: "51467d22d8cc72929e2a028a0428f2086b092bb31cfb79c69cc0a90ce135fde4"
+          },
+          {
+            url: "http://codh.rois.ac.jp/kmnist/dataset/kmnist/train-labels-idx1-ubyte.gz",
+            sha256: "e38f9ebcd0f3ebcdec7fc8eabdcdaef93bb0df8ea12bee65224341c8183d8e17"
+          },
+          {
+            url: "http://codh.rois.ac.jp/kmnist/dataset/kmnist/t10k-images-idx3-ubyte.gz",
+            sha256: "edd7a857845ad6bb1d0ba43fe7e794d164fe2dce499a1694695a792adfac43c5"
+          },
+          {
+            url: "http://codh.rois.ac.jp/kmnist/dataset/kmnist/t10k-labels-idx1-ubyte.gz",
+            sha256: "20bb9a0ef54c7db3efc55a92eef5582c109615df22683c380526788f98e42a1c"
+          }
+        ]
+      end
+    end
+  end
+end

data/lib/torchvision/datasets/mnist.rb

@@ -1,31 +1,10 @@
 module TorchVision
   module Datasets
-    class MNIST
-      RESOURCES = [
-        {
-          url: "http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
-          sha256: "440fcabf73cc546fa21475e81ea370265605f56be210a4024d2ca8f203523609"
-        },
-        {
-          url: "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
-          sha256: "3552534a0a558bbed6aed32b30c495cca23d567ec52cac8be1a0730e8010255c"
-        },
-        {
-          url: "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
-          sha256: "8d422c7b0a1c1c79245a5bcf07fe86e33eeafee792b84584aec276f5a2dbc4e6"
-        },
-        {
-          url: "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz",
-          sha256: "f7ae60f92e00ec6debd23a6088c31dbd2371eca3ffa0defaefb259924204aec6"
-        }
-      ]
-      TRAINING_FILE = "training.pt"
-      TEST_FILE = "test.pt"
-
-      def initialize(root, train: true, download: false, transform: nil)
-        @root = root
+    class MNIST < VisionDataset
+      # http://yann.lecun.com/exdb/mnist/
+      def initialize(root, train: true, download: false, transform: nil, target_transform: nil)
+        super(root, transform: transform, target_transform: target_transform)
         @train = train
-        @transform = transform
 
         self.download if download
 
@@ -33,34 +12,37 @@ module TorchVision
           raise Error, "Dataset not found. You can use download: true to download it"
         end
 
-        data_file = @train ? TRAINING_FILE : TEST_FILE
+        data_file = @train ? training_file : test_file
         @data, @targets = Torch.load(File.join(processed_folder, data_file))
       end
 
       def size
-        @data.size[0]
+        @data.size(0)
       end
 
       def [](index)
-        img = @data[index]
+        img, target = @data[index], @targets[index].item
+
+        img = Utils.image_from_array(img)
+
         img = @transform.call(img) if @transform
 
-        target = @targets[index].item
+        target = @target_transform.call(target) if @target_transform
 
         [img, target]
       end
 
       def raw_folder
-        File.join(@root, "MNIST", "raw")
+        File.join(@root, self.class.name.split("::").last, "raw")
       end
 
       def processed_folder
-        File.join(@root, "MNIST", "processed")
+        File.join(@root, self.class.name.split("::").last, "processed")
       end
 
       def check_exists
-        File.exist?(File.join(processed_folder, TRAINING_FILE)) &&
-          File.exist?(File.join(processed_folder, TEST_FILE))
+        File.exist?(File.join(processed_folder, training_file)) &&
+          File.exist?(File.join(processed_folder, test_file))
       end
 
       def download
@@ -69,7 +51,7 @@ module TorchVision
         FileUtils.mkdir_p(raw_folder)
         FileUtils.mkdir_p(processed_folder)
 
-        RESOURCES.each do |resource|
+        resources.each do |resource|
           filename = resource[:url].split("/").last
           download_file(resource[:url], download_root: raw_folder, filename: filename, sha256: resource[:sha256])
         end
@@ -85,14 +67,43 @@ module TorchVision
           unpack_mnist("t10k-labels-idx1-ubyte", 8, [10000])
         ]
 
-        Torch.save(training_set, File.join(processed_folder, TRAINING_FILE))
-        Torch.save(test_set, File.join(processed_folder, TEST_FILE))
+        Torch.save(training_set, File.join(processed_folder, training_file))
+        Torch.save(test_set, File.join(processed_folder, test_file))
 
         puts "Done!"
       end
 
       private
 
+      def resources
+        [
+          {
+            url: "http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
+            sha256: "440fcabf73cc546fa21475e81ea370265605f56be210a4024d2ca8f203523609"
+          },
+          {
+            url: "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
+            sha256: "3552534a0a558bbed6aed32b30c495cca23d567ec52cac8be1a0730e8010255c"
+          },
+          {
+            url: "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
+            sha256: "8d422c7b0a1c1c79245a5bcf07fe86e33eeafee792b84584aec276f5a2dbc4e6"
+          },
+          {
+            url: "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz",
+            sha256: "f7ae60f92e00ec6debd23a6088c31dbd2371eca3ffa0defaefb259924204aec6"
+          }
+        ]
+      end
+
+      def training_file
+        "training.pt"
+      end
+
+      def test_file
+        "test.pt"
+      end
+
       def unpack_mnist(path, offset, shape)
         path = File.join(raw_folder, "#{path}.gz")
         File.open(path, "rb") do |f|
@@ -101,45 +112,6 @@ module TorchVision
           Torch.tensor(Numo::UInt8.from_string(gz.read, shape))
         end
       end
-
-      def download_file(url, download_root:, filename:, sha256:)
-        FileUtils.mkdir_p(download_root)
-
-        dest = File.join(download_root, filename)
-        return dest if File.exist?(dest)
-
-        temp_path = "#{Dir.tmpdir}/#{Time.now.to_f}" # TODO better name
-
-        digest = Digest::SHA256.new
-
-        uri = URI(url)
-
-        # Net::HTTP automatically adds Accept-Encoding for compression
-        # of response bodies and automatically decompresses gzip
-        # and deflateresponses unless a Range header was sent.
-        # https://ruby-doc.org/stdlib-2.6.4/libdoc/net/http/rdoc/Net/HTTP.html
-        Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == "https") do |http|
-          request = Net::HTTP::Get.new(uri)
-
-          puts "Downloading #{url}..."
-          File.open(temp_path, "wb") do |f|
-            http.request(request) do |response|
-              response.read_body do |chunk|
-                f.write(chunk)
-                digest.update(chunk)
-              end
-            end
-          end
-        end
-
-        if digest.hexdigest != sha256
-          raise Error, "Bad hash: #{digest.hexdigest}"
-        end
-
-        FileUtils.mv(temp_path, dest)
-
-        dest
-      end
     end
   end
 end

data/lib/torchvision/datasets/vision_dataset.rb

@@ -0,0 +1,67 @@
+module TorchVision
+  module Datasets
+    class VisionDataset < Torch::Utils::Data::Dataset
+      attr_reader :data, :targets
+
+      def initialize(root, transforms: nil, transform: nil, target_transform: nil)
+        @root = root
+
+        has_transforms = !transforms.nil?
+        has_separate_transform = !transform.nil? || !target_transform.nil?
+        if has_transforms && has_separate_transform
+          raise ArgumentError, "Only transforms or transform/target_transform can be passed as argument"
+        end
+
+        @transform = transform
+        @target_transform = target_transform
+
+        if has_separate_transform
+          # transforms = StandardTransform.new(transform, target_transform)
+        end
+        @transforms = transforms
+      end
+
+      private
+
+      def download_file(url, download_root:, filename:, sha256:)
+        FileUtils.mkdir_p(download_root)
+
+        dest = File.join(download_root, filename)
+        return dest if File.exist?(dest)
+
+        temp_path = "#{Dir.tmpdir}/#{Time.now.to_f}" # TODO better name
+
+        uri = URI(url)
+
+        # Net::HTTP automatically adds Accept-Encoding for compression
+        # of response bodies and automatically decompresses gzip
+        # and deflateresponses unless a Range header was sent.
+        # https://ruby-doc.org/stdlib-2.6.4/libdoc/net/http/rdoc/Net/HTTP.html
+        Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == "https") do |http|
+          request = Net::HTTP::Get.new(uri)
+
+          puts "Downloading #{url}..."
+          File.open(temp_path, "wb") do |f|
+            http.request(request) do |response|
+              response.read_body do |chunk|
+                f.write(chunk)
+              end
+            end
+          end
+        end
+
+        unless check_integrity(temp_path, sha256)
+          raise Error, "Bad hash"
+        end
+
+        FileUtils.mv(temp_path, dest)
+
+        dest
+      end
+
+      def check_integrity(path, sha256)
+        File.exist?(path) && Digest::SHA256.file(path).hexdigest == sha256
+      end
+    end
+  end
+end

data/lib/torchvision/models/alexnet.rb

@@ -0,0 +1,42 @@
+module TorchVision
+  module Models
+    class AlexNet < Torch::NN::Module
+      def initialize(num_classes: 1000)
+        super()
+        @features = Torch::NN::Sequential.new(
+          Torch::NN::Conv2d.new(3, 64, 11, stride: 4, padding: 2),
+          Torch::NN::ReLU.new(inplace: true),
+          Torch::NN::MaxPool2d.new(3, stride: 2),
+          Torch::NN::Conv2d.new(64, 192, 5, padding: 2),
+          Torch::NN::ReLU.new(inplace: true),
+          Torch::NN::MaxPool2d.new(3, stride: 2),
+          Torch::NN::Conv2d.new(192, 384, 3, padding: 1),
+          Torch::NN::ReLU.new(inplace: true),
+          Torch::NN::Conv2d.new(384, 256, 3, padding: 1),
+          Torch::NN::ReLU.new(inplace: true),
+          Torch::NN::Conv2d.new(256, 256, 3, padding: 1),
+          Torch::NN::ReLU.new(inplace: true),
+          Torch::NN::MaxPool2d.new(3, stride: 2),
+        )
+        @avgpool = Torch::NN::AdaptiveAvgPool2d.new([6, 6])
+        @classifier = Torch::NN::Sequential.new(
+          Torch::NN::Dropout.new,
+          Torch::NN::Linear.new(256 * 6 * 6, 4096),
+          Torch::NN::ReLU.new(inplace: true),
+          Torch::NN::Dropout.new,
+          Torch::NN::Linear.new(4096, 4096),
+          Torch::NN::ReLU.new(inplace: true),
+          Torch::NN::Linear.new(4096, num_classes)
+        )
+      end
+
+      def forward(x)
+        x = @features.call(x)
+        x = @avgpool.call(x)
+        x = Torch.flatten(x, 1)
+        x = @classifier.call(x)
+        x
+      end
+    end
+  end
+end