torchvision 0.1.1 → 0.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c6ae301007e3310dfc378a7f76d24bd71cc048be6156b32d9e0705412817c565
-  data.tar.gz: 423aee54316683ddf9d3893e04f8836dd16ade1408f6c046f6d8f6f337a09a46
+  metadata.gz: 107e429f990a063e57f6218ee6d4fbed3cff80aa6746f0794798d59e1c13b099
+  data.tar.gz: 556fdc4c413803d415ea5575747ec2239f18b0dff9b39a37a4fd6366adec37a6
 SHA512:
-  metadata.gz: cf9c5a514c2ef42299161f97f077897c11b4708858018ffefcb282d6434003d49555ba54ea93ef0fb9e2f3cafdd7bd4395677e02bdfa5b1dcad4772aaeccc896
-  data.tar.gz: 0e01311267e620be6dfbed77724baed42ad9c597a2b5452d762845597ae257932e76848e0069149598743d78dc08b0374b18897ce9a272640eed4a3cec21bbd9
+  metadata.gz: 21b712578516c146888be30bed64a6da6339a42974f5c88fa685278caa9231e4bc4b75e250af3ad37d5450cd205891b31b281cf23fa362456c7ae00998c2736d
+  data.tar.gz: ac86f13e8b5d6a400842ba37b1bf593139360a8df9a19ab4674e71a4327c821a6d5735774185dcbaed368ee735b00bf33d9faedb13d46cd168314da89d900c38

data/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.1.2 (2020-04-29)
+
+- Added CIFAR10, CIFAR100, FashionMNIST, and KMNIST datasets
+- Added ResNet18 model
+
 ## 0.1.1 (2020-04-28)
 
 - Removed `mini_magick` for performance

data/README.md

@@ -20,11 +20,33 @@ This library follows the [Python API](https://pytorch.org/docs/master/torchvisio
 
 ## Datasets
 
-MNIST dataset
+Load a dataset
 
 ```ruby
-trainset = TorchVision::Datasets::MNIST.new("./data", train: true, download: true)
-trainset.size
+TorchVision::Datasets::MNIST.new("./data", train: true, download: true)
+```
+
+Supported datasets are:
+
+- CIFAR10
+- CIFAR100
+- FashionMNIST
+- KMNIST
+- MNIST
+
+## Transforms
+
+```ruby
+TorchVision::Transforms::Compose.new([
+  TorchVision::Transforms::ToTensor.new,
+  TorchVision::Transforms::Normalize.new([0.1307], [0.3081])
+])
+```
+
+## Models
+
+```ruby
+TorchVision::Models::Resnet18.new
 ```
 
 ## Disclaimer

data/lib/torchvision.rb

@@ -6,13 +6,25 @@ require "torch"
 require "digest"
 require "fileutils"
 require "net/http"
+require "rubygems/package"
 require "tmpdir"
 
 # modules
 require "torchvision/version"
 
 # datasets
+require "torchvision/datasets/vision_dataset"
+require "torchvision/datasets/cifar10"
+require "torchvision/datasets/cifar100"
 require "torchvision/datasets/mnist"
+require "torchvision/datasets/fashion_mnist"
+require "torchvision/datasets/kmnist"
+
+# models
+require "torchvision/models/basic_block"
+require "torchvision/models/bottleneck"
+require "torchvision/models/resnet"
+require "torchvision/models/resnet18"
 
 # transforms
 require "torchvision/transforms/compose"

data/lib/torchvision/datasets/cifar10.rb

@@ -0,0 +1,116 @@
+module TorchVision
+  module Datasets
+    class CIFAR10 < VisionDataset
+      # https://www.cs.toronto.edu/~kriz/cifar.html
+
+      def initialize(root, train: true, download: false, transform: nil, target_transform: nil)
+        super(root, transform: transform, target_transform: target_transform)
+        @train = train
+
+        self.download if download
+
+        if !_check_integrity
+          raise Error, "Dataset not found or corrupted. You can use download=True to download it"
+        end
+
+        downloaded_list = @train ? train_list : test_list
+
+        @data = String.new
+        @targets = String.new
+
+        downloaded_list.each do |file|
+          file_path = File.join(@root, base_folder, file[:filename])
+          File.open(file_path, "rb") do |f|
+            while !f.eof?
+              f.read(1) if multiple_labels?
+              @targets << f.read(1)
+              @data << f.read(3072)
+            end
+          end
+        end
+
+        @targets = @targets.unpack("C*")
+        # TODO switch i to -1 when Numo supports it
+        @data = Numo::UInt8.from_binary(@data).reshape(@targets.size, 3, 32, 32)
+        @data = @data.transpose(0, 2, 3, 1)
+      end
+
+      def size
+        @data.shape[0]
+      end
+
+      def [](index)
+        # TODO remove trues when Numo supports it
+        img, target = @data[index, true, true, true], @targets[index]
+
+        # TODO convert to image
+        img = @transform.call(img) if @transform
+
+        target = @target_transform.call(target) if @target_transform
+
+        [img, target]
+      end
+
+      def _check_integrity
+        root = @root
+        (train_list + test_list).each do |fentry|
+          fpath = File.join(root, base_folder, fentry[:filename])
+          return false unless check_integrity(fpath, fentry[:sha256])
+        end
+        true
+      end
+
+      def download
+        if _check_integrity
+          puts "Files already downloaded and verified"
+          return
+        end
+
+        download_file(url, download_root: @root, filename: filename, sha256: tgz_sha256)
+
+        path = File.join(@root, filename)
+        File.open(path, "rb") do |io|
+          Gem::Package.new("").extract_tar_gz(io, @root)
+        end
+      end
+
+      private
+
+      def base_folder
+        "cifar-10-batches-bin"
+      end
+
+      def url
+        "https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz"
+      end
+
+      def filename
+        "cifar-10-binary.tar.gz"
+      end
+
+      def tgz_sha256
+        "c4a38c50a1bc5f3a1c5537f2155ab9d68f9f25eb1ed8d9ddda3db29a59bca1dd"
+      end
+
+      def train_list
+        [
+          {filename: "data_batch_1.bin", sha256: "cee916563c9f80d84e3cc88e17fdc0941787f1244f00a67874d45b261883ada5"},
+          {filename: "data_batch_2.bin", sha256: "a591ca11fa1708a91ee40f54b3da4784ccd871ecf2137de63f51ada8b3fa57ed"},
+          {filename: "data_batch_3.bin", sha256: "bbe8596564c0f86427f876058170b84dac6670ddf06d79402899d93ceea26f67"},
+          {filename: "data_batch_4.bin", sha256: "014e562d6e23c72197cc727519169a60359f5eccd8945ad5a09d710285ff4e48"},
+          {filename: "data_batch_5.bin", sha256: "755304fc0b379caeae8c14f0dac912fbc7d6cd469eb67a1029a08a39453a9add"},
+        ]
+      end
+
+      def test_list
+        [
+          {filename: "test_batch.bin", sha256: "8e2eb146ae340b09e24670f29cabc6326dba54da8789dab6768acf480273f65b"}
+        ]
+      end
+
+      def multiple_labels?
+        false
+      end
+    end
+  end
+end

data/lib/torchvision/datasets/cifar100.rb

@@ -0,0 +1,41 @@
+module TorchVision
+  module Datasets
+    class CIFAR100 < CIFAR10
+      # https://www.cs.toronto.edu/~kriz/cifar.html
+
+      private
+
+      def base_folder
+        "cifar-100-binary"
+      end
+
+      def url
+        "https://www.cs.toronto.edu/~kriz/cifar-100-binary.tar.gz"
+      end
+
+      def filename
+        "cifar-100-binary.tar.gz"
+      end
+
+      def tgz_sha256
+        "58a81ae192c23a4be8b1804d68e518ed807d710a4eb253b1f2a199162a40d8ec"
+      end
+
+      def train_list
+        [
+          {filename: "train.bin", sha256: "f31298fc616915fa142368359df1c4ca2ae984d6915ca468b998a5ec6aeebf29"}
+        ]
+      end
+
+      def test_list
+        [
+          {filename: "test.bin", sha256: "d8b1e6b7b3bee4020055f0699b111f60b1af1e262aeb93a0b659061746f8224a"}
+        ]
+      end
+
+      def multiple_labels?
+        true
+      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/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,11 @@
 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 +13,36 @@ 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
+
+        # TODO convert to image
         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,66 @@
+module TorchVision
+  module Datasets
+    # TODO inherit Torch::Utils::Data::Dataset
+    class VisionDataset
+      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/basic_block.rb

@@ -0,0 +1,46 @@
+module TorchVision
+  module Models
+    class BasicBlock < Torch::NN::Module
+      def initialize(inplanes, planes, stride: 1, downsample: nil, groups: 1, base_width: 64, dilation: 1, norm_layer: nil)
+        super()
+        norm_layer ||= Torch::NN::BatchNorm2d
+        if groups != 1 || base_width != 64
+          raise ArgumentError, "BasicBlock only supports groups=1 and base_width=64"
+        end
+        if dilation > 1
+          raise NotImplementedError, "Dilation > 1 not supported in BasicBlock"
+        end
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        @conv1 = Torch::NN::Conv2d.new(inplanes, planes, 3, stride: stride, padding: 1, groups: 1, bias: false, dilation: 1)
+        @bn1 = norm_layer.new(planes)
+        @relu = Torch::NN::ReLU.new(inplace: true)
+        @conv2 = Torch::NN::Conv2d.new(planes, planes, 3, stride: 1, padding: 1, groups: 1, bias: false, dilation: 1)
+        @bn2 = norm_layer.new(planes)
+        @downsample = downsample
+        @stride = stride
+      end
+
+      def forward(x)
+        identity = x
+
+        out = @conv1.call(x)
+        out = @bn1.call(out)
+        out = @relu.call(out)
+
+        out = @conv2.call(out)
+        out = @bn2.call(out)
+
+        identity = @downsample.call(x) if @downsample
+
+        out += identity
+        out = @relu.call(out)
+
+        out
+      end
+
+      def self.expansion
+        1
+      end
+    end
+  end
+end

data/lib/torchvision/models/bottleneck.rb

@@ -0,0 +1,47 @@
+module TorchVision
+  module Models
+    class Bottleneck < Torch::NN::Module
+      def initialize(inplanes, planes, stride: 1, downsample: nil, groups: 1, base_width: 64, dilation: 1, norm_layer: nil)
+        super()
+        norm_layer ||= Torch::NN::BatchNorm2d
+        width = (planes * (base_width / 64.0)).to_i * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        @conv1 = Torch::NN::Conv2d.new(inplanes, width, 1, stride: 1, bias: false)
+        @bn1 = norm_layer.new(width)
+        @conv2 = Torch::NN::Conv2d.new(width, width, 3, stride: stride, padding: dilation, groups: groups, bias: false, dilation: dilation)
+        @bn2 = norm_layer.new(width)
+        @conv3 = Torch::NN::Conv2d.new(width, planes * self.class.expansion, 1, stride: 1, bias: false)
+        @bn3 = norm_layer.new(planes * self.class.expansion)
+        @relu = Torch::NN::ReLU.new(inplace: true)
+        @downsample = downsample
+        @stride = stride
+      end
+
+      def forward(x)
+        identity = x
+
+        out = @conv1.call(x)
+        out = @bn1.call(out)
+        out = @relu.call(out)
+
+        out = @conv2.call(out)
+        out = @bn2.call(out)
+        out = @relu.call(out)
+
+        out = @conv3.call(out)
+        out = @bn3.call(out)
+
+        identity = @downsample.call(x) if @downsample
+
+        out += identity
+        out = @relu.call(out)
+
+        out
+      end
+
+      def self.expansion
+        4
+      end
+    end
+  end
+end

data/lib/torchvision/models/resnet.rb

@@ -0,0 +1,107 @@
+module TorchVision
+  module Models
+    class ResNet < Torch::NN::Module
+      def initialize(block, layers, num_classes=1000, zero_init_residual: false,
+        groups: 1, width_per_group: 64, replace_stride_with_dilation: nil, norm_layer: nil)
+
+        super()
+        norm_layer ||= Torch::NN::BatchNorm2d
+        @norm_layer = norm_layer
+
+        @inplanes = 64
+        @dilation = 1
+        if replace_stride_with_dilation.nil?
+          # each element in the tuple indicates if we should replace
+          # the 2x2 stride with a dilated convolution instead
+          replace_stride_with_dilation = [false, false, false]
+        end
+        if replace_stride_with_dilation.length != 3
+          raise ArgumentError, "replace_stride_with_dilation should be nil or a 3-element tuple, got #{replace_stride_with_dilation}"
+        end
+        @groups = groups
+        @base_width = width_per_group
+        @conv1 = Torch::NN::Conv2d.new(3, @inplanes, 7, stride: 2, padding: 3, bias: false)
+        @bn1 = norm_layer.new(@inplanes)
+        @relu = Torch::NN::ReLU.new(inplace: true)
+        @maxpool = Torch::NN::MaxPool2d.new(3, stride: 2, padding: 1)
+        @layer1 = _make_layer(block, 64, layers[0])
+        @layer2 = _make_layer(block, 128, layers[1], stride: 2, dilate: replace_stride_with_dilation[0])
+        @layer3 = _make_layer(block, 256, layers[2], stride: 2, dilate: replace_stride_with_dilation[1])
+        @layer4 = _make_layer(block, 512, layers[3], stride: 2, dilate: replace_stride_with_dilation[2])
+        @avgpool = Torch::NN::AdaptiveAvgPool2d.new([1, 1])
+        @fc = Torch::NN::Linear.new(512 * block.expansion, num_classes)
+
+        modules.each do |m|
+          case m
+          when Torch::NN::Conv2d
+            Torch::NN::Init.kaiming_normal!(m.weight, mode: "fan_out", nonlinearity: "relu")
+          when Torch::NN::BatchNorm2d, Torch::NN::GroupNorm
+            Torch::NN::Init.constant!(m.weight, 1)
+            Torch::NN::Init.constant!(m.bias, 0)
+          end
+        end
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual
+          modules.each do |m|
+            case m
+            when Bottleneck
+              Torch::NN::Init.constant!(m.bn3.weight, 0)
+            when BasicBlock
+              Torch::NN::Init.constant!(m.bn2.weight, 0)
+            end
+          end
+        end
+      end
+
+      def _make_layer(block, planes, blocks, stride: 1, dilate: false)
+        norm_layer = @norm_layer
+        downsample = nil
+        previous_dilation = @dilation
+        if dilate
+          @dilation *= stride
+          stride = 1
+        end
+        if stride != 1 || @inplanes != planes * block.expansion
+          downsample = Torch::NN::Sequential.new(
+            Torch::NN::Conv2d.new(@inplanes, planes * block.expansion, 1, stride: stride, bias: false),
+            norm_layer.new(planes * block.expansion)
+          )
+        end
+
+        layers = []
+        layers << block.new(@inplanes, planes, stride: stride, downsample: downsample, groups: @groups, base_width: @base_width, dilation: previous_dilation, norm_layer: norm_layer)
+        @inplanes = planes * block.expansion
+        (blocks - 1).times do
+          layers << block.new(@inplanes, planes, groups: @groups, base_width: @base_width, dilation: @dilation, norm_layer: norm_layer)
+        end
+
+        Torch::NN::Sequential.new(*layers)
+      end
+
+      def _forward_impl(x)
+        x = @conv1.call(x)
+        x = @bn1.call(x)
+        x = @relu.call(x)
+        x = @maxpool.call(x)
+
+        x = @layer1.call(x)
+        x = @layer2.call(x)
+        x = @layer3.call(x)
+        x = @layer4.call(x)
+
+        x = @avgpool.call(x)
+        x = Torch.flatten(x, 1)
+        x = @fc.call(x)
+
+        x
+      end
+
+      def forward(x)
+        _forward_impl(x)
+      end
+    end
+  end
+end

data/lib/torchvision/models/resnet18.rb

@@ -0,0 +1,15 @@
+module TorchVision
+  module Models
+    module ResNet18
+      def self.new(pretrained: false, **kwargs)
+        model = ResNet.new(BasicBlock, [2, 2, 2, 2], **kwargs)
+        if pretrained
+          url = "https://download.pytorch.org/models/resnet18-5c106cde.pth"
+          state_dict = Torch::Hub.load_state_dict_from_url(url)
+          model.load_state_dict(state_dict)
+        end
+        model
+      end
+    end
+  end
+end

data/lib/torchvision/transforms/functional.rb

@@ -22,18 +22,35 @@ module TorchVision
           if std.to_a.any? { |v| v == 0 }
             raise ArgumentError, "std evaluated to zero after conversion to #{dtype}, leading to division by zero."
           end
-          # if mean.ndim == 1
-          #   raise Torch::NotImplementedYet
-          # end
-          # if std.ndim == 1
-          #   raise Torch::NotImplementedYet
-          # end
+          if mean.ndim == 1
+            mean = mean[0...mean.size(0), nil, nil]
+          end
+          if std.ndim == 1
+            std = std[0...std.size(0), nil, nil]
+          end
           tensor.sub!(mean).div!(std)
           tensor
         end
 
         # TODO improve
         def to_tensor(pic)
+          if !pic.is_a?(Numo::NArray) && !pic.is_a?(Torch::Tensor)
+            raise ArgumentError, "pic should be tensor or Numo::NArray. Got #{pic.class.name}"
+          end
+
+          if pic.is_a?(Numo::NArray) && ![2, 3].include?(pic.ndim)
+            raise ArgumentError, "pic should be 2/3 dimensional. Got #{pic.dim} dimensions."
+          end
+
+          if pic.is_a?(Numo::NArray)
+            if pic.ndim == 2
+              raise Torch::NotImplementedYet
+            end
+
+            img = Torch.from_numo(pic.transpose(2, 0, 1))
+            return img.float.div(255)
+          end
+
           pic = pic.float
           pic.unsqueeze!(0).div!(255)
         end

data/lib/torchvision/version.rb

@@ -1,3 +1,3 @@
 module TorchVision
-  VERSION = "0.1.1"
+  VERSION = "0.1.2"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: torchvision
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.1.2
 platform: ruby
 authors:
 - Andrew Kane
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2020-04-28 00:00:00.000000000 Z
+date: 2020-04-29 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: numo-narray
@@ -30,14 +30,14 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 0.2.3
+        version: 0.2.4
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 0.2.3
+        version: 0.2.4
 - !ruby/object:Gem::Dependency
   name: bundler
   requirement: !ruby/object:Gem::Requirement
@@ -90,7 +90,16 @@ files:
 - LICENSE.txt
 - README.md
 - lib/torchvision.rb
+- lib/torchvision/datasets/cifar10.rb
+- lib/torchvision/datasets/cifar100.rb
+- lib/torchvision/datasets/fashion_mnist.rb
+- lib/torchvision/datasets/kmnist.rb
 - lib/torchvision/datasets/mnist.rb
+- lib/torchvision/datasets/vision_dataset.rb
+- lib/torchvision/models/basic_block.rb
+- lib/torchvision/models/bottleneck.rb
+- lib/torchvision/models/resnet.rb
+- lib/torchvision/models/resnet18.rb
 - lib/torchvision/transforms/compose.rb
 - lib/torchvision/transforms/functional.rb
 - lib/torchvision/transforms/normalize.rb