这是一篇面向新手的博文：因为本人也是新手，记录一下自己在做这个项目遇到的大大小小的坑。
按照下面的例子写就好了
import torch as t
from torch.utils import data
import os
from PIL import Image
import numpy as np
from torchvision import transforms as T
from torch import nn
from torch.autograd import Variable
from torch.optim import Adam
from torchvision.utils import make_grid
from torch.utils.data import DataLoader transform = T.Compose( #归一化
[
T.Resize(1000),
T.CenterCrop(1000),
T.ToTensor(),
T.Normalize(mean=[.5,.5,.5],std=[.5,.5,.5])
]
) #在这个class里重写你的数据集
class twoface(data.Dataset):
def __init__(self,root,transforms=None):
imgs = os.listdir(root)
self.imgs = [os.path.join(root,img) for img in imgs]
self.transforms = transforms def __getitem__(self, index):
img_path = self.imgs[index]
label = 1 if 'empty' in img_path.split('/')[-1] else 0 #定义标签：图片名中有label
data = Image.open(img_path)
if self.transforms:
data = self.transforms(data)
return data,label def __len__(self):
return len(self.imgs) dataset = twoface('./data1/',transforms=transform) #transform 用在这里
data_loader_train = t.utils.data.DataLoader(dataset=dataset,
batch_size=1, #batch_size 要=1
shuffle=True,
num_workers=0) data_loader_test = t.utils.data.DataLoader(dataset=dataset,
batch_size=1,
shuffle=True,
num_workers=0) class Model(t.nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv1 = t.nn.Sequential(t.nn.Conv2d(3, 6, kernel_size=20, stride=10, padding=0), #模型这块如果你图片格式不对的话要重新算
t.nn.ReLU(),
t.nn.Conv2d(6, 10, kernel_size=6, stride=1, padding=0),
t.nn.ReLU(),
t.nn.Conv2d(10, 16, kernel_size=5, stride=1, padding=0),
t.nn.ReLU(),
t.nn.MaxPool2d(stride=5, kernel_size=5))
#
self.dense = t.nn.Sequential(t.nn.Linear(18 * 18 * 16, 33),
t.nn.ReLU(),
t.nn.Linear(33, 2)
) def forward(self, x):
x = self.conv1(x)
x = x.view(-1,18 * 18 * 16)
x = self.dense(x)
return x model = Model().cuda()
print(model) cost = t.nn.CrossEntropyLoss()
optimizer = t.optim.Adam(model.parameters())
n_epochs = 5 for epoch in range(n_epochs):
running_loss = 0.0
running_correct = 0
print("Epoch {}/{}".format(epoch, n_epochs))
print("-" * 10) for datas in data_loader_train:
realone,label = datas
realone, label = Variable(realone).cuda(), Variable(label).cuda()
output = model(realone)
_,pred = t.max(output.data,1)
optimizer.zero_grad()
loss = cost(output, label)
loss.backward()
optimizer.step()
running_loss += loss.data[0]
running_correct += t.sum(pred == label) testing_correct = 0
for datak in data_loader_test:
X_test, y_test = datak
X_test, y_test = Variable(X_test).cuda(), Variable(y_test).cuda()
outputs = model(X_test)
_, pred = t.max(outputs.data, 1)
testing_correct += t.sum(pred == y_test.data)
print(
"Loss is:{:.4f}, Train Accuracy is:{:.4f}%, Test Accuracy is:{:.4f}".format(running_loss / len(dataset),
100 * running_correct / len(
dataset),
100 * testing_correct / len(
dataset)))
t.save(model, 'ifempty.pkl')
运行model中的坑
from pylab import plt
import torch as t
from torch.autograd import Variable
from torchvision.utils import make_grid
from PIL import Image
import numpy as np
import os class Model(t.nn.Module): #如果不把这些乱七八糟的class和transform 重写一遍会出错 def __init__(self):
super(Model, self).__init__()
self.conv1 = t.nn.Sequential(t.nn.Conv2d(3, 6, kernel_size=20, stride=10, padding=0),
t.nn.ReLU(),
t.nn.Conv2d(6, 10, kernel_size=6, stride=1, padding=0),
t.nn.ReLU(),
t.nn.Conv2d(10, 16, kernel_size=5, stride=1, padding=0),
t.nn.ReLU(),
t.nn.MaxPool2d(stride=5, kernel_size=5))
#
self.dense = t.nn.Sequential(t.nn.Linear(18 * 18 * 16, 33),
t.nn.ReLU(),
t.nn.Linear(33, 2)
) def forward(self, x):
x = self.conv1(x)
x = x.view(-1,18 * 18 * 16)
x = self.dense(x)
return x lst = os.listdir('./recv/')
flag = 0
for i in range(len(lst)): img=Image.open('./recv/' + lst[i]) from torchvision import transforms as T
trans= T.Compose(
[
T.Resize(1000),
T.CenterCrop(1000),
T.ToTensor(),
T.Normalize(mean=[.5,.5,.5],std=[.5,.5,.5])
]
) a=trans(img)
b = a.numpy()
x = np.array([b])
y = t.Tensor(x)
fix_noise = Variable(y)
fix_noise = fix_noise.cuda()
net = t.load('ifempty.pkl')
output = net(fix_noise)
_,pred = t.max(output.data,1) if pred.data.item() == 1:
print('空')
flag = 1
else:
flag = 0 if flag == 0:
netj = t.load('ifzwh.pkl')
output2 = netj(fix_noise)
_, pred2 = t.max(output2.data, 1)
if pred2.data.item() == 0:
print('xxx')
else:
print('其他人')