【写给自己】自定义自己的数据并调用Torch自带的ResNet网络
很多时候可能会用到别人已经搭建好的网络，而恰好pytorch内置许多优秀的网络，在实际应用中已经可以满足需求了。
数据预处理
此时数据不再是one_folder_one_label，本次用的是来自kaggle的dogs-vs-cats。是一个二分类问题。目录如下
其中train数据是“cat.1.jpg”形式，在导入数据时候要提取出cat关键词，并将其划分label。并且因为不能用ImageFolder直接导入数据了，因此需要完成继承Dataset，再自定义自己的数据。本次代码同样来自https://www.bilibili.com/video/BV1354y1s7kQ
class Mydata(Dataset):
def __init__(self, file_path, trans):
self.file_path = file_path
self.images = os.listdir(self.file_path) # 把路径下所有文件放在一起
self.transform = trans
def __len__(self):
return len(self.images)
def __getitem__(self, index): # 根据index 返回img和label
image_index = self.images[index]
img_path = os.path.join(self.file_path, image_index)
img = Image.open(img_path).convert('RGB')
# 原图片命名规则:cat.2.jpg，得到cat
label = img_path.split('\\')[-1].split('.')[0]
label = 1 if 'dog' in label else 0
if self.transform is not None:
img = self.transform(img)
return img, label
定义完数据类后就可以创建数据了
# Data
transform = transforms.Compose([
transforms.CenterCrop(200),
transforms.ToTensor(),
transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
])
dataset = Mydata(path, transform)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
调用已有的ResNet网络
需要将pretrained部分设为False，用自己的数据来训练
# Model
device = torch.device("cuda:0" if torch.cuda.is_available() == 1 else "cpu")
# 调用已有网络
net = torchvision.models.resnet18(pretrained=False) # False，用自己的数据来训练
# 全连接层特征数
num_labels = net.fc.in_features
# 修改特征数为2
net.fc = nn.Linear(num_labels, 2)
net.to(device)
同时定义训练代码为
def train(model, train_data, epoch, device):
running_loss = 0
model.train()
for x, y in train_data:
x, y = x.to(device), y.to(device)
y_pred = model(x)
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss = loss_func(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
with torch.no_grad():
running_loss += loss.item()
epoch_loss = running_loss / len(train_data.dataset)
print('epoch: ', epoch + 1,
'loss: ', round(epoch_loss, 3))
return epoch_loss
之后进行训练，训练完后仅仅保存网络的参数，而不是完整的网络
print("Start Training")
# Train
train_loss = []
for epoch in range(epochs):
epoch_loss = train(net, dataloader, epoch, device)
train_loss.append(epoch_loss)
torch.save(net.state_dict(), 'resnet.pt')
训练时间大概是100分钟
训练部分完整代码
import torch
import torch.nn as nn
import torchvision
import numpy as np
from torch.utils.data import Dataset, DataLoader
import os
from PIL import Image
from torchvision import transforms
import torch.nn.functional as F
import time
class Mydata(Dataset):
def __init__(self, file_path, trans):
self.file_path = file_path
self.images = os.listdir(self.file_path) # 把路径下所有文件放在一起
self.transform = trans
def __len__(self):
return len(self.images)
def __getitem__(self, index): # 根据index 返回img和label
image_index = self.images[index]
img_path = os.path.join(self.file_path, image_index)
img = Image.open(img_path).convert('RGB')
# 原图片命名规则:cat.2.jpg，得到cat
label = img_path.split('\\')[-1].split('.')[0]
label = 1 if 'dog' in label else 0
if self.transform is not None:
img = self.transform(img)
return img, label
def train(model, train_data, epoch, device):
running_loss = 0
model.train()
for x, y in train_data:
x, y = x.to(device), y.to(device)
y_pred = model(x)
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss = loss_func(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
with torch.no_grad():
running_loss += loss.item()
epoch_loss = running_loss / len(train_data.dataset)
print('epoch: ', epoch + 1,
'loss: ', round(epoch_loss, 3))
return epoch_loss
if __name__ == '__main__':
start = time.time()
path = './dogs_vs_cats/train/train'
epochs = 30
# Data
transform = transforms.Compose([
transforms.CenterCrop(200),
transforms.ToTensor(),
transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
])
dataset = Mydata(path, transform)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
# Model
device = torch.device("cuda:0" if torch.cuda.is_available() == 1 else "cpu")
# 调用已有网络
net = torchvision.models.resnet18(pretrained=False) # False，用自己的数据来训练
# 全连接层特征数
num_labels = net.fc.in_features
# 修改特征数为2
net.fc = nn.Linear(num_labels, 2)
net.to(device)
print("Start Training")
# Train
train_loss = []
for epoch in range(epochs):
epoch_loss = train(net, dataloader, epoch, device)
train_loss.append(epoch_loss)
torch.save(net.state_dict(), 'resnet.pt')
end = time.time()
print(end - start)
调用训练好的网络来测试
对于测试集依然需要继承Dataset来得到自己的数据。训练时候发现，虽然labels是list数据，但是list里面的数据是tensor类型，用torch.save()保存后可能会是乱码，于是先将list里数据转化为numpy，再用numpy.save()保存。但不知道效果如何。似乎测试出来的labels很多不匹配，不知道该如何验证结果。
import numpy as np
import torch
from torchvision import transforms
from torch.utils.data import Dataset
import torch.nn as nn
import torchvision
import os
from PIL import Image
import torch.nn.functional as F
class Testdata(Dataset):
def __init__(self, path, transform):
self.path = path
self.image = os.listdir(self.path)
self.trans = transform
def __len__(self):
return len(self.image)
def __getitem__(self, index):
image_index = self.image[index]
img_path = os.path.join(self.path, image_index)
img = Image.open(img_path).convert('RGB')
if self.trans is not None:
img = self.trans(img)
return img
def classifier(net, data, device):
labels_ = []
for x in data:
x = x.to(device)
x = x.view(1, 3, 200, 200)
out_ = net(x)
out_ = F.softmax(out_, dim=1)
label_ = torch.max(out_, 1).indices
labels_.append(label_)
return labels_
if __name__ == "__main__":
# 加载网络
device = torch.device("cuda:0" if torch.cuda.is_available() == 1 else "cpu")
net = torchvision.models.resnet18(pretrained=False)
num_labels = net.fc.in_features
net.fc = nn.Linear(num_labels, 2)
net.to(device)
# 将网络参数调整后再load参数
state_dict = torch.load('resnet.pt')
net.load_state_dict(state_dict)
net.eval()
# 加载数据
path = './dogs_vs_cats/test1/test1'
transform = transforms.Compose([
transforms.CenterCrop(200),
transforms.ToTensor(),
transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
])
unlabeled_data = Testdata(path, transform)
# 测试
labels = classifier(net, unlabeled_data, device)
label = []
for i in range(len(labels)):
label.append(labels[i].cpu().numpy())
print('done')
np.save('label.csv', label )