I successfully exported my model to a .pkl file. I am deploying my model on a Jetson nano GPU. The problem is that I created a custom transform (RandFlip) that I included in the item_tfms parameter of DataBlock. I want this transform to be applied only in training and not in inference. My input is a camera image. On the nano I am feeding the camera image directly to the predict function without going through a data loader. I do not want the transform RandFlip to be applied in deploying the model.
The problem I have is that the export operation fails complaining that it does not recognize my RandFlip class. This should not be a problem because RandFlip should not be applied in deployment. If I remove the RandFlip class in training, the export to .pkl works fine and I can do prediction on the nano.
So my question, is there a way to do augmentation with my custom RandFlip class, export my model to a .pkl file, and do inference bypassing augmentation?
Create a new dummy class RandFlip that just does nothing but return the inputs as it is. That being said Iâm also looking forward to better answers on how to handle these cases efficiently.
Setting split_idx allowed the export to work successfully. But in loading âexport.pklâ on the nano I now get this error.
AttributeError: Canât get attribute âRandFlipâ on <module âmainâ from âtestdrive.pyâ>
I can probably solve this by adding the RandFlip class to my testdrive.py but this should be unnecessary.
That is correct. Because all the fastai stuff is inside the namespace available when you load it in. Your custom class is not. It exports references to it, but not the actual class itself. The fastai bits work the exact same way as well, you have fastai installed it can grab those references (and the steps to recreate your Learner from it)
The same thing happens with the acc_camvid function for anyone trying to export and bring in a segmentation model from the course (itâs a very very often asked question on the forums)
I can concur Iâm still in Undergrad Just come with a general idea of the course will not prepare you for every scenario, and over time you will learn much more Python (which is why he recommends coding for at least 1 year before taking the course).
To followup, this is the transform I want to happen in training but not inference. I thought that since this is called in the DataLoader it would not be called in inference because I am using learn.predict(image) which presumably bypasses the DataLoader. So, setting both split_idx=1 and p=0 do not prevent the transform being applied in inference. The transform happens about half the time like it was in training with p=.5. This makes no sense to me.
class RandFlip(RandTransform):
def __init__(self, p=0, **kwargs):
super().__init__(p=p, split_idx=1, **kwargs)
def encodes(self, o: Tensor):
return Tensor([-o[0], o[1]]) # this reverses the target steering value
def encodes(self, o: PILImage):
print("This should not happen") # with p=0 this should not print
return ImageOps.mirror(o) # this mirrors the input image
Does calling an instance of it and checking what itâs split idx show 1?
mytfm = RandFlip()
mytfm.split_idx
Otherwise I would set it up similar to the other transforms like so:
class RandFlip(RandTransform):
split_idx=1
def __init__(self, p=0, **kwargs):
super().__init__(p=p, **kwargs)
def encodes(self, o: Tensor):
return Tensor([-o[0], o[1]]) # this reverses the target steering value
def encodes(self, o: PILImage):
print("This should not happen") # with p=0 this should not print
return ImageOps.mirror(o) # this mirrors the input image
Update. I solved my problem by not exporting to a .pkl file. Instead I saved the state dict on the training side which turns out saving a lot of space (pickling: 120M vs state_dict: 1.6M).
torch.save(model.state_dict(), new_model_path)
Then on the inference side I just included the model definition and loaded an image like this. Who knew it could be so easy.
from torchvision import transforms
from torch.autograd import Variable
model.load_state_dict(torch.load(new_model_path)
def image_loader(image_path):
"""load image, returns cuda tensor"""
image = Image.open(image_path)
image = transforms.ToTensor()(image).float()
image = Variable(image, requires_grad=True)
image = image.unsqueeze(0)
#return image.cuda() #assumes that you're using GPU
return image #assumes no GPU
image = image_loader(image_path)
model(image) # this does the prediction
Iâm still in Undergrad 