mirror of
https://github.com/wassname/multifit.git
synced 2026-09-11 12:20:41 +08:00
36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
from typing import Iterator, Collection
|
|
from fastai.data_block import CategoryListBase
|
|
from fastai.text import *
|
|
|
|
|
|
class BinaryProcessor(CategoryProcessor):
|
|
def create_classes(self, classes):
|
|
self.classes = classes
|
|
if classes is not None: self.c2i = {0:0, 1:1}
|
|
def generate_classes(self, items):
|
|
return [0]
|
|
|
|
class BinaryCategoryList(CategoryListBase):
|
|
"Basic `ItemList` for single classification labels."
|
|
_processor=BinaryProcessor
|
|
def __init__(self, items:Iterator, classes:Collection=None, label_delim:str=None, **kwargs):
|
|
super().__init__(items, classes=classes, **kwargs)
|
|
mean = self.items.mean()
|
|
# if mean and mean != 0:
|
|
# weight = torch.tensor([1 / mean]).cuda()
|
|
# print(f'Weighting BCEWithLogitsFlat by {weight.item()}')
|
|
# else:
|
|
weight = None
|
|
# raise Exception('debug')
|
|
self.loss_func = BCEWithLogitsFlat(weight=weight)
|
|
|
|
def reconstruct(self, t):
|
|
return Category(t, self.c2i[t.item()])
|
|
|
|
def get(self, i):
|
|
o = self.items[i]
|
|
if o is None: return None
|
|
return Category(o, self.c2i[o])
|
|
|
|
def analyze_pred(self, pred, thresh:float=0.5): return pred.argmax()
|