Torchvision框架学习之FCOS模型及其训练

from torchvision.models import detectionmodel = detection.fcos_resnet50_fpn(progress=True,num_classes=3,          pretrained_backbone=True,          trainable_backbone_layers=4)

https://aistudio.baidu.com/aistudio/datasetdetail/6045

{'labels': tensor([1, 1, 1, 2, 2, 2, 2, 2]),  'boxes': tensor([[ 711,  233,  844,  506],          [1036,  194, 1206,  459],          [ 958,  406, 1239,  573],          [1142,  194, 1275,  320],          [ 780,  478,  908,  614],          [ 766,  612,  914,  742],          [ 972,  542, 1120,  678],          [ 986,  684, 1120,  820]])}

from pathlib import Pathfrom torchvision.io import read_image,ImageReadModeimport jsonimport torchclass BNDataset(torch.utils.data.Dataset):    def __init__(self, istrain=True,datapath='D:/data/lslm'):   #注意修改数据集路径        self.datadir=Path(datapath)/('train' if istrain else 'test')        self.idxfile=self.datadir/('train.txt' if istrain else 'test.txt')        self.labelnames=['background','bolt','nut']        self.data=self.parseidxfile()    def parseidxfile(self):        lines=open(self.idxfile).readlines()        return [line for line in lines if len(line)>5]

    def __getitem__(self, idx):        data=self.data[idx].split('\t')        x = read_image(str(self.datadir/data[0]),ImageReadMode.RGB)/255.0        labels=[]        boxes=[]        for i in data[2:]:            if len(i)<5: continue            r=json.loads(i)            labels.append(self.labelnames.index( r['value']))            cords=r['coordinate']            xyxy=cords[0][0],cords[0][1],cords[1][0],cords[1][1]            boxes.append(xyxy)        y = {            'labels': torch.LongTensor(labels),            'boxes': torch.tensor(boxes).long()        }        return x, y

    def __len__(self):        return len(self.data)

def collate_fn(data):    x = [i[0] for i in data]    y = [i[1] for i in data]    return x, y

train_loader = torch.utils.data.DataLoader(dataset=BNDataset(istrain=True), batch_size=4, shuffle=True, drop_last=True, collate_fn=collate_fn)test_loader = torch.utils.data.DataLoader(dataset=BNDataset(istrain=False), batch_size=1, shuffle=True, drop_last=True, collate_fn=collate_fn)

for i, (x, y) in enumerate(train_loader):    labels=[loader.dataset.labelnames[i] for i in y[0]['labels']]colors=[ ('red' if i=='nut' else 'blue') for i in labels]image=draw_bounding_boxes(x[0], y[0]['boxes'],labels=labels,colors=colors,width=5,font_size=50,outtype='CHW')vis.image(image)

def train():    model.train()    optimizer = torch.optim.SGD(model.parameters(), lr=0.0001,momentum=0.98)    for epoch in range(5):        for i, (x, y) in enumerate(train_loader):            outs = model(x, y)            loss = outs['classification']+ outs['bbox_ctrness']+outs['bbox_regression']            loss.backward()            optimizer.step()            optimizer.zero_grad()            if i % 10 == 0:                print(epoch, i, loss.item())        torch.save(model, f'./models/tvs{epoch}.model')
train()

0 0 2.1866644620895386......4 100 0.9854792356491089

def test():    model_load = torch.load('./models/tvs4.model')    model_load.eval()    loader_test = torch.utils.data.DataLoader(dataset=BNDataset(istrain=False), batch_size=1, shuffle=False, drop_last=True, collate_fn=collate_fn)

    for i, (x, y) in enumerate(loader_test):        with torch.no_grad():            outs = model_load(x)            res=outs[0]            boxes=res['boxes']            scores=res['scores']            labels=res['labels']            #阈值过滤            threhold=0.5   #保留类别概率大于0.5的检测结果            mask=scores>threhold            scores=scores[mask]            labels=labels[mask]            boxes=boxes[mask]            labelnames=[loader_test.dataset.labelnames[i]+f'{scores[idx]:.2f}' for idx,i in enumerate(labels)]            colors=[ ('red' if i==1 else 'blue') for i in labels]            img=draw_bounding_boxes(x[0],boxes,labels=labelnames,colors=colors,width=3,font_size=50,outtype=’CHW’)vis.image(img)