import os
import numpy as np
import random
import cv2
import torch
import torch.nn as nn
import torch.utils.data as dataf
from torch.autograd import Variable
import torch.nn.functional as F
import torch.optim as optim
train_image_path = r"F:test\image"
signal = os.listdir(train_image_path)
all_path = []
data_path = train_image_path
for fsingal in signal:
filepath = data_path+"/"+fsingal
filename = os.listdir(filepath)
for fname in filename:
ffpath = filepath+"/"+fname
path = [fsingal,ffpath]
all_path.append(path)
count = len(all_path)
data_x = np.empty((count,1,28,28),dtype="float32")
data_y = []
random.shuffle(all_path)
i=0
for item in all_path:
img = cv2.imread(item[1],0)
img = cv2.resize(img,(28,28))
arr = np.asarray(img,dtype="float32")
data_x[i ,:,:,:] = arr
i+=1
data_y.append(int(item[0]))
data_x = data_x / 255
data_y = np.asarray(data_y,)
data_x = torch.from_numpy(data_x)
data_y = torch.from_numpy(data_y).long() # 将数据类型转为Long
#制作数据集
dataset = dataf.TensorDataset(data_x,data_y)
#加载数据,为之后的训练准备
loader = dataf.DataLoader(dataset, batch_size=batchsize, shuffle=True)
class L5_NET(nn.Module):
def __init__(self):
super(L5_NET,self).__init__();
self.conv1 = nn.Conv2d(1,20,kernel_size=5)
self.conv2 = nn.Conv2d(20,30,kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(30*4*4,50)
self.fc2 = nn.Linear(50,2)
def forward(self,x):
x = F.relu(F.max_pool2d(self.conv1(x),2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)),2))
x = x.view(-1,30*4*4)
x = F.relu(self.fc1(x))
x = F.dropout(x,p=0.8,training=self.training)
x = self.fc2(x)
return x
model = L5_NET()
cuda = False
train_epoch = 20
train_lr = 0.01
train_momentum = 0.5
batchsize = 3
optimizer = optim.SGD(model.parameters(),lr = train_lr,momentum = train_momentum)
criterion = nn.CrossEntropyLoss()
def train(epoch):
model.train()
for batch_idx, (data, target) in enumerate(loader):
if cuda:
data, target = data.cuda(), target.cuda()
data, target = Variable(data), Variable(target)
#求导
optimizer.zero_grad()
#训练模型,输出结果
output = model(data)
#在数据集上预测loss
loss = F.nll_loss(output, target)
#反向传播调整参数
loss.backward()
#SGD刷新进步
optimizer.step()
#实时输出
if batch_idx % 10 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(loader.dataset),100. * batch_idx / len(loader), loss.data[0]))
在这期间有报错Assertion `cur_target >= 0 && cur_target < n_classes’ failed,在loss = F.nll_loss(output, target),原因是因为我的数据集类别从1开始,那么,类别数就要加1。