Press "Enter" to skip to content

CNN–MNIST识别和自己数据集

本小节使用torch搭建CNN模型,训练和测试:

(1)定义模型超参数:输出、迭代次数、批量大小、学习率。

(2)定义训练数据,加餐部分是使用自己的数据集:(可参考:https://blog.csdn.net/u014365862/article/details/80506147)

(3)定义模型(定义卷积神经网络)。

(4)定义损失函数,选用适合回归问题的损失函数。

(5)定义优化算法(SGD、Adam等)。

(6)保存模型。

———————————我是可爱的分割线———————————

代码部分:

# coding=utf-8import torch import torch.nn as nnimport torchvisionimport torchvision.transforms as transforms# 判定GPU是否存在device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')# 定义超参数num_epochs = 5num_classes = 10batch_size = 100learning_rate = 0.001# 手写体数据,(数据+标签)train_dataset = torchvision.datasets.MNIST(root='./data/',                                           train=True,                                            transform=transforms.ToTensor(),                                           download=True)test_dataset = torchvision.datasets.MNIST(root='./data/',                                          train=False,                                           transform=transforms.ToTensor())# # 构建数据管道, 使用自己的数据集请参考:https://blog.csdn.net/u014365862/article/details/80506147train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)# 定义自己的卷积神经网络class ConvNet(nn.Module):
def __init__(self, num_classes=10):
super(ConvNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.layer2 = nn.Sequential(
nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.layer3 = nn.Sequential(
nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.fc = nn.Linear(3*3*64, num_classes)

def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = self.layer3(out)
out = out.reshape(out.size(0), -1)
out = self.fc(out)
return out

# 定义模型
model = ConvNet(num_classes).to(device)

# 定义损失函数+优化算法
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
total_step = len(train_loader)
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)

# 前向传播+计算loss
outputs = model(images)
loss = criterion(outputs, labels)

# 后向传播+调整参数
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 每100个batch打印一次数据
if (i+1) % 100 == 0:
print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
.format(epoch+1, num_epochs, i+1, total_step, loss.item()))

# 模型测试部分
# 测试阶段不需要计算梯度,注意
model.eval() # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()

print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))

# 保存模型参数
torch.save(model.state_dict(), 'model.ckpt')

加餐:在自己数据集上使用:

其中,train.txt中的数据格式:

gender/0male/0(2).jpg 1

gender/0male/0(3).jpeg 1

gender/0male/0(1).jpg 0

test.txt中的数据格式如下:

gender/0male/0(3).jpeg 1

gender/0male/0(1).jpg 0

gender/1female/1(6).jpg 1

代码部分:

# coding=utf-8import torch import torch.nn as nnimport torchvisionfrom torch.utils.data import Dataset, DataLoader    from torchvision import transforms, utils from PIL import Image # 判定GPU是否存在device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')# 定义超参数num_epochs = 5num_classes = 10batch_size = 16learning_rate = 0.001def default_loader(path):        # 注意要保证每个batch的tensor大小时候一样的。    return Image.open(path).convert('RGB')        class MyDataset(Dataset):        def __init__(self, txt, transform=None, target_transform=None, loader=default_loader):            fh = open(txt, 'r')            imgs = []            for line in fh:                line = line.strip('n')                # line = line.rstrip()            words = line.split(' ')                imgs.append((words[0],int(words[1])))            self.imgs = imgs            self.transform = transform            self.target_transform = target_transform            self.loader = loader                def __getitem__(self, index):            fn, label = self.imgs[index]            img = self.loader(fn)            if self.transform is not None:                img = self.transform(img)            return img,label                def __len__(self):            return len(self.imgs)        def get_loader(dataset='train.txt', crop_size=128, image_size=28, batch_size=2, mode='train', num_workers=1):        """Build and return a data loader."""        transform = []        if mode == 'train':            transform.append(transforms.RandomHorizontalFlip())        transform.append(transforms.CenterCrop(crop_size))        transform.append(transforms.Resize(image_size))        transform.append(transforms.ToTensor())        transform.append(transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))        transform = transforms.Compose(transform)        train_data=MyDataset(txt=dataset, transform=transform)        data_loader = DataLoader(dataset=train_data,                                      batch_size=batch_size,                                      shuffle=(mode=='train'),                                      num_workers=num_workers)        return data_loader    # 注意要保证每个batch的tensor大小时候一样的。# data_loader = DataLoader(train_data, batch_size=2,shuffle=True)train_loader = get_loader('train.txt', batch_size=batch_size)    print(len(train_loader))    test_loader = get_loader('test.txt', batch_size=batch_size)    print(len(test_loader))    # 定义自己的卷积神经网络class ConvNet(nn.Module):    def __init__(self, num_classes=10):        super(ConvNet, self).__init__()        self.layer1 = nn.Sequential(            nn.Conv2d(3, 16, kernel_size=5, stride=1, padding=2),            nn.BatchNorm2d(16),
           nn.ReLU(),            nn.MaxPool2d(kernel_size=2, stride=2))        self.layer2 = nn.Sequential(            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),            nn.BatchNorm2d(32),            nn.ReLU(),            nn.MaxPool2d(kernel_size=2, stride=2))        self.layer3 = nn.Sequential(            nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),            nn.BatchNorm2d(64),            nn.ReLU(),            nn.MaxPool2d(kernel_size=2, stride=2))        self.fc = nn.Linear(3*3*64, num_classes)            def forward(self, x):        out = self.layer1(x)        out = self.layer2(out)        out = self.layer3(out)        out = out.reshape(out.size(0), -1)        out = self.fc(out)        return out# 定义模型model = ConvNet(num_classes).to(device)# 定义损失函数+优化算法criterion = nn.CrossEntropyLoss()optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)# 训练模型total_step = len(train_loader)for epoch in range(num_epochs):    for i, (images, labels) in enumerate(train_loader):        images = images.to(device)        labels = labels.to(device)                # 前向传播+计算loss        outputs = model(images)        loss = criterion(outputs, labels)                # 后向传播+调整参数        optimizer.zero_grad()        loss.backward()        optimizer.step()        # 每100个batch打印一次数据        if (i+1) % 100 == 0:            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'                    .format(epoch+1, num_epochs, i+1, total_step, loss.item()))# 模型测试部分# 测试阶段不需要计算梯度,注意model.eval()  # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)with torch.no_grad():    correct = 0    total = 0    for images, labels in test_loader:        images = images.to(device)        labels = labels.to(device)        outputs = model(images)        _, predicted = torch.max(outputs.data, 1)        total += labels.size(0)        correct += (predicted == labels).sum().item()    print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))# 保存模型参数torch.save(model.state_dict(), 'model.ckpt')

总结:

本节使用CNN训练MNIST识别、自己数据的识别。

上面加餐部分需要生成自己的txt文件(数据+标签),可以参考这个,自己以前调试用的:https://github.com/MachineLP/py_workSpace/blob/master/g_img_path.py

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注