数据准备

1.第一种方式

torchvision.datasets.ImageFolder


数据集文件需如下结构:

rootdir/cat/xxx.png

root_dir/cat/xxy.jpeg



root_dir/dog/nsdf3.png

root_dir/dog/asd932
.png

import torch
import torchvision
from torchvision import transforms,datasets
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
transform = transforms.Compose([
    transforms.Resize([224,224]),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = datasets.ImageFolder(root="/Volumes/storage/learn/machine-learning/notes/pytorch/data/dog_cat",
                               transform=transform)
dataloader = torch.utils.data.DataLoader(dataset,batch_size=4,shuffle=True,num_workers=2)
import matplotlib.pyplot as plt
import numpy as np

label_dict = ("cat","dog")
# 输出图像的函数
def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 随机获取训练图片
dataiter = iter(dataloader)
images, labels = dataiter.next()

# 显示图片
imshow(torchvision.utils.make_grid(images))
plt.show()
print(' '.join('%5s' % label_dict[labels[j].item()] for j in range(4)))

png

dog cat dog cat


第二种方式(继承torch.utils.data.Dataset)

实现:

__len__(self): 返回数据集的大小

__getitem__(self,index): 支持索引,获取第index个样本的数据和标签


假设数据集文件如下结构:

root_dir/cat.0.png

root_dir/cat.1.jpeg



root_dir/dog.0.jpg

root_dir/dog.1.png



note:transform的输入图片格式:PIL.Image

import torch
import os
from PIL import Image
from skimage import io
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
labels = ("cat","dog")
class dogCatData(torch.utils.data.Dataset):
    def __init__(self,root_dir,transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.files = []
        for file in os.listdir(root_dir):
            sample = {}
            sample["file"] = os.path.join(root_dir,file)
            sample["label"] = labels.index(file.split(".")[0])
            self.files.append(sample)
    def __len__(self):
        return len(self.files )
    def __getitem__(self,index):
        sample = self.files[index]
        #image = Image.open(sample["file"])
        image = io.imread(sample["file"])
        image = Image.fromarray(image)
        label = sample["label"]
        if self.transform is not None:
            image = self.transform(image)
        return image,label
data = dogCatData("/Volumes/storage/learn/bigdata/data/dog_cat/train")
print(data[0])
(<PIL.Image.Image image mode=RGB size=208x257 at 0x1492917F0>, 0)
import torchvision
from torchvision import transforms as transforms
transform = transforms.Compose([
    transforms.Resize([224,224]),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = dogCatData("/Volumes/storage/learn/bigdata/data/dog_cat/train",transform)
dataloader = torch.utils.data.DataLoader(dataset,batch_size=8,shuffle=True,num_workers=4)
import matplotlib.pyplot as plt
import numpy as np

label_dict = ("cat","dog")
# 输出图像的函数
def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 随机获取训练图片
dataiter = iter(dataloader)
images, labels = dataiter.next()

# 显示图片
imshow(torchvision.utils.make_grid(images))
plt.show()
print(' '.join('%5s' % label_dict[labels[j].item()] for j in range(8)))

png

cat cat dog dog dog dog dog cat


构建模型

import torch
import torch.nn as nn
class AlexNet(nn.Module):
    def __init__(self,num_classes=2):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x
model = AlexNet(2)

定义损失和优化器

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

模型训练

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)
model.to(device)
model.to(device)
for epoch in range(4):  # loop over the dataset multiple times
    model.train()
    running_loss = 0.0
    for i, data in enumerate(dataloader):
        # get the inputs
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % 1999 == 0:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')
torch.save(model, 'dog_cat.pth')
[1,     1] loss: 0.636
[1,     2] loss: 1.332
[1,     3] loss: 2.026
[1,     4] loss: 2.692
[1,     5] loss: 3.386
[1,     6] loss: 4.139
[1,     7] loss: 4.776
[1,     8] loss: 5.472
[1,     9] loss: 6.108
[1,    10] loss: 6.860
[1,    11] loss: 7.555
[1,    12] loss: 8.281
[1,    13] loss: 8.915
[1,    14] loss: 9.641
[1,    15] loss: 10.338
[1,    16] loss: 11.062
[1,    17] loss: 11.697
Finished Training

模型保存和加载

1.仅保存模型参数

import torch
torch.save(model.state_dict(), path) # save
model.load_state_dict(torch.load(path)) # load
# save
state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
torch.save(state, path)
# load
checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint(['epoch'])

2.保存/加载整个模型

torch.save(model, path)
model = torch.load(path)


深度学习      pytorch

本博客所有文章除特别声明外,均采用 CC BY-SA 3.0协议 。转载请注明出处!