猫狗大战——pytorch+resnet18

数据

这是一个在kaggle上的竞赛，原数据提供了25000张图片，本文所使用的数据集来自其中train的8000张,包括4000张猫和4000张狗

#将数据导入
data_dir = 'E:\\code\\Python\\catanddog\\train'
test_dir = 'E:\\code\\Python\\catanddog\\test'

class Data(data.Dataset):
    def __init__(self, path, transform = None, train = True, test = False):
        self.test = test
        self.train = train
        self.transform = transform
        #imgs = [os.path.join(path, img) for img in os.listdir(path)]
        imgs = [os.path.join(data_dir, img) for img in path]#imgs存的是每张图片的总路径
            
        if self.test:#test模式
            self.imgs = imgs
        else:#train模式
            random.shuffle(imgs)#数据打乱
            self.imgs = imgs
            
    def __getitem__(self, index):
        img = self.imgs[index]
        if self.test:
            label = 2
        else:
            label = 0 if 'cat' in img.split('\\')[-1] else 1#cat = 0, dog = 1
        image = Image.open(img)
        image = self.transform(image)
        return image, label
    
    def __len__(self):
        return len(self.imgs)

#train和val的transform处理
#对train进行了随机裁剪翻转的操作
transform_train = transforms.Compose([
    transforms.Resize((256, 256)), 
    transforms.RandomCrop((224, 224)), 
    transforms.RandomHorizontalFlip(), 
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.2225)) 
])


transform_val = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 
train_dirs, valid_dirs = train_test_split(os.listdir(data_dir), test_size = 0.2, random_state = 2021)
trainset = Data(train_dirs, transform = transform_train)
valset = Data(valid_dirs, transform = transform_val)
trainloader = torch.utils.data.DataLoader(trainset, batch_size = 20, shuffle = True, num_workers = 0)
valloader = torch.utils.data.DataLoader(valset, batch_size = 20, shuffle = False, num_workers = 0)

Net

model = resnet18(pretrained=True)
model.fc = nn.Linear(512, 2,bias=True)
model = model.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)  # 设置训练细节
scheduler = StepLR(optimizer, step_size=3)
criterion = nn.CrossEntropyLoss()

train+test

#训练并保存模型
for epoch in range(1):
    train(epoch)
    val(epoch)
torch.save(model, 'catvsdog_model.pth')  # 保存模型
#一次epoch在验证集上准确度有0.988

#对图片进行预测
classes = ['cat', 'dog']
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.load('catvsdog_model.pth')  # 加载模型
model = model.to(device)
model.eval()  # 把模型转为test模式

img = cv2.imread("E:\\code\\Python\\catanddog\\test\\dog.3886.jpg")  # 读取要预测的图片
cv2.imshow("img", img)
img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
trans = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])

img = trans(img)
img = img.to(device)
img = img.unsqueeze(0)  # 图片扩展多一维,因为输入到保存的模型中是4维的[batch_size,通道,长，宽]，而普通图片只有三维，[通道,长，宽]
output = model(img)
prob = F.softmax(output, dim=1)  # prob是2个分类的概率
value, predicted = torch.max(output.data, 1)
pred_class = classes[predicted.item()]
print(pred_class)