PyTorch implementation of neural networks for solving problems from the Machine Number Sense (MNS) dataset [1]. Dataset and official implementation of baseline models can be found in this repo, created by paper authors.
$ pip install machine_number_sense
MLP [1]:
import torch
from mns.model import ConvMLP
x = torch.rand(4, 3, 80, 80)
mlp = ConvMLP(image_size=80)
logits = mlp(x)
logits # torch.Tensor with shape (4, 99)
LSTM [1]:
import torch
from mns.model import ConvLSTM
x = torch.rand(4, 3, 80, 80)
lstm = ConvLSTM(image_size=80)
logits = lstm(x)
logits # torch.Tensor with shape (4, 99)
Scattering Compositional Learner (SCL) [2] adapted to problems from the MNS dataset:
import torch
from mns.model import SCL
x = torch.rand(4, 3, 80, 80)
scl = SCL(image_size=80)
logits = scl(x)
logits # torch.Tensor with shape (4, 99)
Implementation of SCL for solving Raven's Progressive Matrices can be found in this repo.
Neural Arithmetic Logic Unit (NALU) [3] adapted to MNS:
import torch
from mns.model import ConvNALU
x = torch.rand(4, 3, 80, 80)
nalu = ConvNALU(image_size=80)
logits = nalu(x)
logits # torch.Tensor with shape (4, 99)
The MNS dataset can be obtained as described in this repo. After downloading, it can be loaded with:
from mns.dataset import MNSDataset
dataset = MNSDataset(data_dir='/path/to/dataset', image_size=80)
iterator = iter(dataset)
image, target = next(iterator)
image # torch.Tensor with shape (3, 80, 80)
target # torch.Tensor with shape ()
File mns.module
contains a PyTorch Lightning module for training models on MNS.
Training can be run with Docker using scripts from the scripts/
directory.
$ python -m pytest tests
[1] Zhang, Wenhe, et al. "Machine number sense: A dataset of visual arithmetic problems for abstract and relational reasoning." Proceedings of the AAAI Conference on Artificial Intelligence. 2020.
[2] Wu, Yuhuai, et al. "The Scattering Compositional Learner: Discovering Objects, Attributes, Relationships in Analogical Reasoning." arXiv preprint arXiv:2007.04212 (2020).
[3] Trask, Andrew, et al. "Neural arithmetic logic units." Advances in Neural Information Processing Systems. 2018.
@inproceedings{zhang2020machine,
title={Machine number sense: A dataset of visual arithmetic problems for abstract and relational reasoning},
author={Zhang, Wenhe and Zhang, Chi and Zhu, Yixin and Zhu, Song-Chun},
booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
volume={34},
number={02},
pages={1332--1340},
year={2020}
}
@article{wu2020scattering,
title={The Scattering Compositional Learner: Discovering Objects, Attributes, Relationships in Analogical Reasoning},
author={Wu, Yuhuai and Dong, Honghua and Grosse, Roger and Ba, Jimmy},
journal={arXiv preprint arXiv:2007.04212},
year={2020}
}
@inproceedings{trask2018neural,
title={Neural arithmetic logic units},
author={Trask, Andrew and Hill, Felix and Reed, Scott E and Rae, Jack and Dyer, Chris and Blunsom, Phil},
booktitle={Advances in Neural Information Processing Systems},
pages={8035--8044},
year={2018}
}