数据准备
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)))
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)))
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)
本博客所有文章除特别声明外,均采用 CC BY-SA 3.0协议 。转载请注明出处!