Re usable Components and CNN Trainer

Hey all, I’m new to Bolts and a bit to lightning, but wokred with PyTorch. This project is super awesome, completely in snyc with my thoughts.

🚀 Feature

Torchivision supports lot of pre-trained CNN models. One of the common tasks in fine tuning these CNNs. (Transfer learning) This issue was open with lightning team as well I guess. Can we create a trainer for this purpose that allows easy fine tuning or even raw training just by using pretrained = False

The only layer that differs in the last Linear layer which being nn.Linear(hidden_dim, num_classes)

Motivation

Such pre-trained CNNs are used as Backbones for Object detection tasks as well. E.g. you can fine tune a CNN on say animals such as tiger, lion and then use these CNN weights as backbone for F-RCNN .

This is a key feature, lightning supports multi GPU training and TPU training out of the box for all such operations.

Pitch

One of my earlier attempts to do the same with pytorch was here. I could successfully tune and train all the PyTorch models and almost all PyTorch Image Models. It had simple API which syncs with #200.

model = model_factory.create_torchvision_model("resnet18", num_classes=10, pretrained=True,)

All these fine-tunable backbones are compatible with F-RCNN also, they can simply be passed to torchvision FRCNN.

Alternatives

I’m unsure if my backbone extraction for CNN is efficient and good. It assumes we have nn.AdaptiveAvgPool2d() followed by Linear only in classification tasks, hence simply extracts all layers before this linear. I held to this assumption since all torchvision and PyTorch Image Models follow this. But I guess we need to assume something.

Additional Context

What this feature allows is simple fine tuner to CNNs directly supported with all lightning perks and bonuses. A common backbone API, re-usable to all vision models such as FRCNN, Mask-RCNN, RetinaNet (in next torchvision release). I’m unsure still if this can be extensible and maintainable.

I can probably work on this feature, I would need a bit help as I’m new to this repo and lightning.

[justusschock]

@oke-aditya please slow down a bit 😄 we’re not yet sure, how/whether this shall be included to bolts at all, so we don’t want anybody to put some effort to this, when this may be discarded later on.

[nateraw]

Yes! I believe this is exactly where we are tryin to go w/ things…

Model architectures that are dead simple to use and include options to:

• load pretrained weights we provide and use them for inference or finetuneing on your task
• load random weights + train yourself
• subclass + do whatever you want

This is not just limited to torchvision CNNs, but more general, I think. Torchvision stuff is a great place to start though. If you want to give implementing this a try, I’d love to see a PR from you 😄 .

---

Feel free to reach out to me on the slack channel and/or make a post on the forum to discuss further.