迁移学习-使用resnet18分类猫狗图片

本分类用的数据集为kaggle上的猫狗图片数据集

首先分成训练集和验证集

在train文件夹下创建dog和cat两个文件夹

train_root = 'D:\kaggle\\train'
dog_folder = os.path.join(train_root,'dog')
cat_folder = os.path.join(train_root,'cat')
if not os.path.exists(dog_folder):
    os.mkdir(dog_folder)
if not os.path.exists(cat_folder):
    os.mkdir(cat_folder)

在val文件夹下创建dog和cat两个文件夹

val_root  = 'D:\kaggle\\val'
dog_folder = os.path.join(val_root, 'dog')
cat_folder = os.path.join(val_root, 'cat')
if not os.path.exists(dog_folder):
    os.mkdir(dog_folder)
if not os.path.exists(cat_folder):
    os.mkdir(cat_folder)


#zip为存放所有数据集的文件夹
data_file = os.listdir('D:\kaggle\\zip')

#分为两类
dog_file  = list(filter(lambda x: x[:3]=='dog', data_file))
cat_file = list(filter(lambda x: x[:3]=='cat', data_file))

#90%作训练集，10%作验证集
for i in range(len(dog_file)):
    pic_path = root + '\\zip\\' + dog_file[i]
    if i<len(dog_file)*0.9:
        obj_path = root + '\\train\\dog\\' + dog_file[i]
    else:
        obj_path = root + '\\val\\dog\\' + dog_file[i]
    shutil.move(pic_path,obj_path)

for i in range(len(cat_file)):
    pic_path = root + '\\zip\\' + cat_file[i]
    if i < len(cat_file) * 0.9:
        obj_path = root + '\\train\\cat\\' + cat_file[i]
    else:
        obj_path = root + '\\val\\cat\\' + cat_file[i]
    shutil.move(pic_path, obj_path)

下面为处理图片数据构造dataloader

data_transforms = {
    'train':
        transforms.Compose([
            transforms.RandomResizedCrop(299),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
        ]),
    'val':
    transforms.Compose([
        transforms.Resize(320),
        transforms.CenterCrop(299),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
}

root = 'D:\kaggle\\'
data_folder = {
    'train':
    ImageFolder(
        os.path.join(root,'train'),transform = data_transforms['train']
    ),
    'val':
    ImageFolder(
        os.path.join(root, 'val'), transform=data_transforms['val']
    )
}

batch_size = 32
dataloader = {
    'train':
        DataLoader(
            data_folder['train'],
            batch_size=batch_size,
            shuffle=True,
            num_workers=0,
        ),
    'val':
        DataLoader(
            data_folder['val'],
            batch_size=batch_size,
            num_workers=0,
        ),
}

data_size = {
    'train':len(dataloader['train'].dataset),
    'val': len(dataloader['val'].dataset)
}

img_classes = len(dataloader['train'].dataset.classes) #2

 

下面使用了resnet18残差网络作迁移学习，

use_gpu = torch.cuda.is_available()
transfer_model = models.resnet18(pretrained=True)
for param in transfer_model.parameters():
    param.require_grad = False          #不改变卷积网络部分的参数
dim_in = transfer_model.fc.in_features  #获取现行层的输入特征维度
transfer_model.fc = nn.Linear(dim_in,2) #改变最后线性层的参数
if use_gpu:
    transfer_model = transfer_model.cuda()

optimizer = optim.Adam(transfer_model.fc.parameters(),lr=1e-3)
criterion = nn.CrossEntropyLoss()
num_epoch = 10

for epoch in range(num_epoch):
    print('{}/{}'.format(epoch + 1, num_epoch))
    print('Train')
    transfer_model.train()
    running_loss = 0.0
    running_acc = 0.0
    since = time.time()
    for i,data in enumerate(dataloader['train'],1):
        img, label = data
        if use_gpu:
            img = Variable(img).cuda()
            label = Variable(label).cuda()

        out = transfer_model(img)
        loss = criterion(out,label)
        _,pred = torch.max(out,1)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_loss += loss.item()*label.size(0)
        acc = torch.sum(pred==label)
        running_acc += acc.item()
        if i % 100 == 0:
            print('Loss: {:.6f}, Acc: {:.4f}'.format(running_loss / (
                    i * batch_size), running_acc / (i * batch_size)))
    running_loss /= data_size['train']
    running_acc /= data_size['train']
    elips_time = time.time() - since

    print('Loss: {:.6f}, Acc: {:.4f}, Time: {:.0f}s'.format(
            running_loss, running_acc, elips_time))

    print('Validation')

    transfer_model.eval()
    num_correct = 0.0
    total = 0.0
    eval_loss = 0.0
    for data in dataloader['val']:
        img,label = data
        with torch.no_grad():
            img = Variable(img).cuda()
        with torch.no_grad():
            label = Variable(label).cuda()
        out = transfer_model(img)
        _, pred = torch.max(out.data, 1)
        loss = criterion(out, label)
        eval_loss += loss.item() * label.size(0)
        num_correct += (pred.cpu() == label.data.cpu()).sum()
        total += label.size(0)
    print('Loss: {:.6f} Acc: {:.4f}'.format(eval_loss / total, num_correct /
                                                total))
    print()
print('Finish Training!')
print()
save_path = os.path.join(root, 'model_save')
if not os.path.exists(save_path):
    os.mkdir(save_path)
torch.save(transfer_model.state_dict(), save_path + '/resnet18.pth')