pytorch训练自己的数据集

    model = models.resnet50(pretrained=False)
    class_num = 62
    fc_features = model.fc.in_features
    model.fc = nn.Linear(fc_features, class_num)
    model = model.to(device)

    import os
    data = []
    labels = []
    filetype = "ppm"
    for root,dirs,files in os.walk("./"):
    for f in files:
        if f.split('.')[1] == filetype:
            data.append(root[1:]+"/"+f)
            labels.append(int(root[2:]))
    
    with open("datalist.txt",'w') as fi:
    for i in range(len(data)-1):
        fi.write("{} {}\n".format(data[i],labels[i]))
    fi.write("{} {}".format(data[i+1],labels[i+1]))

    def default_loader(path):
    return Image.open(path).convert('RGB')
    
    
    class BTSCDataset(Dataset):
    def __init__(self, dir, transforms=None, loader=default_loader, train=True):
        super(BTSCDataset, self).__init__()
        self.sub_directory = 'Training' if train else 'Testing'
        imgs = []
        with open(dir + "/BelgiumTSC_" + self.sub_directory + "/" + self.sub_directory + "/datalist.txt", "r") as f:
            data = f.readlines()
        for line in data:
            label = int(line.split()[1])
            img = line.split()[0]
            imgs.append((dir + "/BelgiumTSC_" + self.sub_directory + "/" + self.sub_directory + img, label))
        self.imgs = imgs
        self.transforms = transforms
        self.loader = loader
    
    def __getitem__(self, item):
        fn, label = self.imgs[item]
        img = self.loader(fn)
        if self.transforms is not None:
            img = self.transforms(self)
        return img, fn
    
    def __len__(self):
        return len(self.imgs)