mauro-nievoff
/

MultiCaReClassifier

Image Classification

Model card Files Files and versions Community

mauro-nievoff commited on 10 days ago

Commit

eed82ca

·

verified ·

1 Parent(s): 04cffe3

Update pipeline.py

Files changed (1) hide show

pipeline.py +1 -1

pipeline.py CHANGED Viewed

@@ -135,7 +135,7 @@ class MultiCaReClassifier():
       checkpoint_file = os.path.join(model_name.replace(':', '_'), 'model')
       dls = ImageDataLoaders.from_path_func('', imgs, lambda x: '0', item_tfms=Resize((224,224), method='squish'))
       learn = vision_learner(dls, resnet50, n_out=len(labels)).to_fp16()
-      learn.load(checkpoint_file, device=device)
       test_dl = learn.dls.test_dl(imgs, device=device)
       probs, _ = learn.get_preds(dl=test_dl)
       self.data.loc[condition, model_name] = labels[probs.argmax(axis=1)]

       checkpoint_file = os.path.join(model_name.replace(':', '_'), 'model')
       dls = ImageDataLoaders.from_path_func('', imgs, lambda x: '0', item_tfms=Resize((224,224), method='squish'))
       learn = vision_learner(dls, resnet50, n_out=len(labels)).to_fp16()
+      learn.load(checkpoint_file, device=device, weights_only=False)
       test_dl = learn.dls.test_dl(imgs, device=device)
       probs, _ = learn.get_preds(dl=test_dl)
       self.data.loc[condition, model_name] = labels[probs.argmax(axis=1)]