I can't create the body of ResNet18 with fastai

Question

I'm trying to build the body of the ResNet18 in this code:

from fastai.vision.data import create_body
from fastai.vision import models
from torchvision.models.resnet import resnet18
from fastai.vision.models.unet import DynamicUnet
import torch

def build_res_unet(n_input=1, n_output=2, size=256):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    body = create_body(resnet18, n_in = n_input, pretrained=True, cut=-2)
    net_G = DynamicUnet(body, n_output, (size, size)).to(device)
    return net_G

net_G = build_res_unet(n_input=1, n_output=2, size=256)

but I keep getting an error:

TypeError: create_body() got an unexpected keyword argument 'n_in'

but in the fastai docs the n_in parameter is present.

How can I create the body, am I missing something?

Answer 1

I tested the code on my local machine and it runs perfectly, maybe there are some problems on Google Colab! I will update this answer if I found a way to make it run on colab

EDIT: I solved the problem by adding !pip install fastai==2.4 on google colab, the version used by colab was very old