Adding Image Dataloaders and Flax Resnet18 model

[shreyaspadhy]

Adding the first combined snippets of code to this repo! This PR contains -

1. Image dataloaders from learning-invariances, with definitions of default transformations.
2. Defining a big METADATA dict in jaxutils.data.image that should contain all the necessary mean, std, size, num_datapoints information for all datasets.
3. Adding a model definition of Resnet18, and a conversion script that converts Pytorch resnet18 models from bayesian-lottery-tickets to Flax.

[JamesAllingham]

Mostly LGTM! Just a few minor suggestions/questions for the most part.

Also, I assume that you didn't mean to upload the __pycache__ folder?

And I think that the img_resnets.py can just be renamed to resnets.py, since a) there are no other kinds of resnets (yet?), and b) by default I think it is assumed that resnets will be the convolutional kind. But that is also minor.

[JamesAllingham]

It isn't really necessary to put this here! As I mentioned, this is code that has been adapted by me and Javi over the course of our PhDs, so this repo isn't really the origin.

[shreyaspadhy]

got it!

[JamesAllingham]

Nice addition!

[JamesAllingham]

This is also good!

[JamesAllingham]

Is this just to say that you need augmentations for this project, or is there something more to this comment? :)

[shreyaspadhy]

I added some more detail. Basically, we need to not use augmented train data, when training for the linear mode, and for sampling from the posterior. This was actually a bug I'd chased down, where if we augment, then the sampling is misspecified, and we never converge.

[JamesAllingham]

For the non-imagenet cases, this makes sense since the test augementations are empty, but for imagenet does it make sense to be applying transformations when the user of the function has set perform_augmentations=False?

[shreyaspadhy]

The reason we need Resize(256) and CenterCrop(224) for Imagenet is because, by default, test images in Imagenet are all of random sizes, and not uniform. So we still need some deterministic preprocessing to ensure all images are of size 224x224x3.

[JamesAllingham]

This docstring isn't exactly accurate, since the augmentations are not limited to affine transformations?

[shreyaspadhy]

good point, affine is wrong there. made the change.

[JamesAllingham]

Maybe this should be a flag that can be set by the user (if it is specific to BLT models)?

[shreyaspadhy]

Will do this when I add the torchhub model conversion fns that you had.

[JamesAllingham]

Since this is very similar to LeNet, with only the sizes changing, I think this is a good place to use inheritance to reduce code duplication. Specifically, I'd change LeNet to also have a separate setup and __call__ and then when defining LeNetSmall inherit from LeNet and simply override the setup.

[shreyaspadhy]

Sounds good! I have a doubt about inheritance. Since I'm redefining setup for LeNetSmall, I'll have to redefine self.dense here, even though it is the same. Is there a way to only overwrite a subset of class properties after inheriting within Flax?

[JamesAllingham]

Hmmm, maybe you can call super().setup() in your inherited class's setup function and then redefine only the attributes which need to change. However, I think this is something I've tried before and not had luck. You'll probably have to bite the bullet and have some code duplication I am afraid.

[JamesAllingham]

To be clear, you can definitely call super.setup() and then add new attributes, however, I think an error is thrown when you try and overwrite attributes which have been defined in the parent's setup.

[JamesAllingham]

OK, I've confirmed that redefining attributes when calling super.setup() inside of the child class's setup() doesn't work. E.g., trying this:

class MyNet(nn.Module):

    def setup(self):
        self.fc2 = nn.Dense(10, 10)
        self.fc1 = nn.Dense(10, 10)

    def __call__(self, x):
        return self.fc2(self.fc1(x))

class MyNet2(MyNet):

    def setup(self):
        super().setup()

        self.fc2 = nn.Dense(333, 10)

    def __call__(self, x):
        return self.fc2(self.fc1(x))

will result in this error, when trying to call init:

ValueError: Duplicate use of scope name: "fc2"

However, if you really want to you can avoid the code duplication like this:

class MyNet(nn.Module):

    def _partial_setup(self):
        self.fc2 = nn.Dense(10, 10)

    def setup(self):
        self._partial_setup()
        self.fc1 = nn.Dense(10, 10)

    def __call__(self, x):
        return self.fc2(self.fc1(x))


class MyNet2(MyNet):

    def _partial_setup(self):
        self.fc2 = nn.Dense(333, 10)

    def __call__(self, x):
        return self.fc2(self.fc1(x))

But I think in this case it is better to have the duplicated code, since it is only 1 line.

[shreyaspadhy]

Thanks for trying these out, this is super useful to know!

[JamesAllingham]

This decorator should be removed if you have a setup method.

[shreyaspadhy]

Is there a way in Flax to mix and match? For example, I define self.conv{i} in the setup, but if I want to compactly define nn.Dense within the call without mentioning it in the setup? I'm curious about the syntax here.

[JamesAllingham]

Currently, it actually doesn't break anything to mix and match setup and compact, however, as far as I understand, this is not advertised functionality and there are plans to remove this functionality in future: google/flax#2018.

Why do you actually want to mux and match though? Seems to me that everything could be in setup here?

[shreyaspadhy]

Ah, good to know! I wanted to be cheeky and define the Dense layer inside, but good to know it's not preferred behaviour.