-
MNIST数据集
我这里使用的数据集为MNIST提供的免费数据集。其中包含6w张训练集和1k张验证集。
在该数据集中,图片都是以pytorch张量的形式表现的。为使方法具有通用性,这里将张量转化为图片并以 data -> train/val -> 类别1/2/... -> 图片的形式进行存储。如下图:
这里的每张照片都单通道28*28的格式。
- 数据整理
为方便调用,我决定生成两个txt文件记录图片信息和对应的标签。如下:
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\1.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\1000.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10005.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10010.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10022.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10025.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10026.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10045.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10069.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10071.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10080.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10083.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10107.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10119.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10120.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10121.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10128.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10140.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\1015.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10167.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10168.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10176.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\1018.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10188.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10195.jpg" 0
"D:\\data\\numbersee\\numbersee\\data\\train\\0\\10203.jpg" 0
为实现这一效果,我们定义一方法
def getdata(inwhile,outwhile):
这里inwhile是data文件夹的位置。outwhile是文件的保存位置。
然后我们提取各个图片的位置,并创建两个字典,以 位置:标签的形式保存到字典当中。
root=os.path.join(os.getcwd(),inwhile)
datalist=os.listdir(root)
train=os.path.join(root,datalist[0])
val =os.path.join(root,datalist[1])
trainlist =os.listdir(train)
vallist =os.listdir(val)
trainout={}
valout ={}
for i in trainlist:
train_i=os.path.join(train,i)
train_i_list=os.listdir(train_i)
for j in train_i_list:
jie=os.path.join(train_i,j)
trainout[jie]=int(i)
for i in vallist:
train_i=os.path.join(val,i)
train_i_list=os.listdir(train_i)
for j in train_i_list:
jie=os.path.join(train_i,j)
valout[jie]=int(i)
最后将字典逐行保存到txt文件中。
train_path=os.path.join(os.getcwd(),outwhile)
train_path=os.path.join(train_path,r'train.txt')
val_path=os.path.join(os.getcwd(),outwhile)
val_path=os.path.join(val_path,r'val.txt')
filet=open(train_path,'w')
filev=open(val_path,'w')
for i in trainout:
jsr=json.dumps(i)
filet.write(jsr)
filet.write(' ')
jsr2=json.dumps(trainout[i])
filet.write(jsr2)
filet.write('\r\n')
for i in valout:
jsv=json.dumps(i)
filev.write(jsv)
filev.write(' ')
jsv2=json.dumps(valout[i])
filev.write(jsv2)
filev.write('\r\n')
filet.close()
filev.close()
全部代码如下:
def getdata(inwhile,outwhile):
root=os.path.join(os.getcwd(),inwhile)
datalist=os.listdir(root)
train=os.path.join(root,datalist[0])
val =os.path.join(root,datalist[1])
trainlist =os.listdir(train)
vallist =os.listdir(val)
trainout={}
valout ={}
for i in trainlist:
train_i=os.path.join(train,i)
train_i_list=os.listdir(train_i)
for j in train_i_list:
jie=os.path.join(train_i,j)
trainout[jie]=int(i)
for i in vallist:
train_i=os.path.join(val,i)
train_i_list=os.listdir(train_i)
for j in train_i_list:
jie=os.path.join(train_i,j)
valout[jie]=int(i)
train_path=os.path.join(os.getcwd(),outwhile)
train_path=os.path.join(train_path,r'train.txt')
val_path=os.path.join(os.getcwd(),outwhile)
val_path=os.path.join(val_path,r'val.txt')
filet=open(train_path,'w')
filev=open(val_path,'w')
for i in trainout:
jsr=json.dumps(i)
filet.write(jsr)
filet.write(' ')
jsr2=json.dumps(trainout[i])
filet.write(jsr2)
filet.write('\r\n')
for i in valout:
jsv=json.dumps(i)
filev.write(jsv)
filev.write(' ')
jsv2=json.dumps(valout[i])
filev.write(jsv2)
filev.write('\r\n')
filet.close()
filev.close()
- 数据预处理及打包
这里我们用到了pytorch中的Dataset类和DataLoader方法。
我们先写一个自己的Dataset类去生成一个数据集。
class Mydate(Dataset):
def __init__(self,date,dev):
self.date=open(date).readlines()
self.dev=dev
def __getitem__(self,index):
out = self.date[index]
label = int(out[-2])
data = out[1:-4]
data = Image.open(data)
to_tensor = transforms.Compose([
transforms.ToTensor(),transforms.Normalize((0.1307,),(0.3081,))])
data = to_tensor(data)
data=data.to(self.dev)
label2 = torch.tensor([0.0]*10)
label2=label2.to(self.dev)
label2[label]=1
return data,label2
def __len__(self):
return len(self.date)
这里可以看到我们在定义一个Mydata类时要传入两个变量一个是包好数据和标签的txt文件,另一个是device类型的参数。这里dev=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
该类每次抛出一个数据时,会将数据正则化。
下面我们再定义一个DataLoader方法。去将数据打包成数据包,并且随机排列一下。
def getbat(data,dev=torch.device("cpu"),size=1,num_workers=0):
mdataset = Mydate(data,dev)
dataloade = DataLoader(dataset=mdataset,batch_size=size,shuffle=True,num_workers=num_workers)
return dataloade
- 模型搭建
数据处理的方法和类完成后,我们接下来进行模型的搭建。我这里使用的网络模型为
卷积层 -> 卷积层 -> 池化层 -> 卷积层 -> 池化层 -> 卷积层 -> 输出层
class number_model(nn.Module):
def __init__(self,):
super(number_model,self).__init__()
self.con1=nn.Conv2d(in_channels=1,out_channels=16,kernel_size=3,stride=1,padding=1)
self.ru1=nn.ReLU()
self.con2=nn.Conv2d(in_channels=16,out_channels=32,kernel_size=3,stride=1,padding=1)
self.ru2=nn.ReLU()
self.pool1=nn.MaxPool2d((2,2),(2,2))
#14*14*32
self.con3=nn.Conv2d(in_channels=32,out_channels=64,kernel_size=3,stride=1,padding=1)
self.ru3=nn.ReLU()
self.pool2=nn.MaxPool2d((2,2),(2,2))
#7*7*64
self.con4=nn.Conv2d(in_channels=64,out_channels=128,kernel_size=3,stride=1,padding=1)
self.ru4=nn.ReLU()
#7*7*128
self.fc=nn.Linear(in_features=(7*7*128),out_features=10)
def forward(self,x):
out=self.con1(x)
out=self.ru1(out)
out=self.con2(out)
out=self.ru2(out)
out=self.pool1(out)
out=self.con3(out)
out=self.ru3(out)
out=self.pool2(out)
out=self.con4(out)
out=self.ru4(out)
out=torch.flatten(out,0)
out=self.fc(out)
return out
- 模型的训练
万事具备只欠东风
这里开始写训练文件
def main():
#调用gpu
dev=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#读取训练集文件夹,并打包成训练集
path=os.path.join(os.getcwd(),r'numberdata')
print(path)
train_path=os.path.join(path,r'train.txt')
val_path=os.path.join(path,r'val.txt')
train_data=getbat(train_path,dev,1)
val_data=getbat(val_path,dev,10)
number_m=number_model()
number_m=number_m.to(dev)
#优化器的调用和模型训练
number_m.train(True)
optimizer = sgd.SGD(number_m.parameters(), lr=1e-6)
loss_fn = nn.CrossEntropyLoss()
n=10
for i in trange(n):
j=1
for x,y_true in tqdm(train_data,'bat'):
y_p = number_m(x[0])
y_true=torch.tensor(y_true[0]).to(dev)
loss = loss_fn(y_p, y_true)
optimizer.zero_grad()
loss.backward()
optimizer.step()
j+=1
if j%1000==0:
print(loss)
#模型的验证
number_m.train(False)
true_num=0
l=[]
for x,y_true in tqdm(val_data,'bat'):
y_p = number_m(x[0])
y_true=torch.tensor(y_true[0]).to(dev)
if int(y_p.argmax().item())==int((y_true.argmax().item())):
true_num+=1
else:
l.append((int(y_p.argmax().item()),int((y_true.argmax().item()))))
#保存
torch.save(number_m, 'number.pth')
#输出识别错误的情况和正确率。
print(l,true_num/len(val_data))
if __name__=='__main__':
main()
经过反复的实验,发现学习率为0.000001,步数为10时。效果最佳。
最终成功率为95%还算符合预期。
- 模型的调用
模型训练完成要使用,这里编写一个方法调用模型。因为模型的输入为单通道28*28,所以使用前要将目的图形转化为灰度图,并且分辨率改为28*28.
def main():
dev=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data=r"C:\Users\meng xiang\Desktop\1.jpg"
data = Image.open(data)
data = data.convert('L')
data = data.resize((28, 28))
data.show()
to_tensor = transforms.Compose([
transforms.ToTensor(),transforms.Normalize((0.1307,),(0.3081,))])
data = to_tensor(data)
print(data.shape)
data=data.to(dev)
model=torch.load('number.pth')
out=model(data)
return '这个图形是:'+str(out.argmax().item())
if __name__=='__main__':
print(main())