pytroch 一个简单例子 一个小demo

Posted 2022-04-08 _刘文凯_

初入pytorch的门，感觉pytorch和tensorflow的思路还是有点区别的，做了一个小的例子，算是开始入门吧 pytorch实现卷积神经完了过 pytroch实现CNN

这里写出来，供大家入门使用

图片数据：

链接：https://pan.baidu.com/s/1fF6wSSj7x19aAi2dhPvCwA
提取码：0sk3

代码：

import torchvision
from torchvision import transforms
from torch.utils import data
from PIL import Image
import torch.nn as nn
import torch.nn.functional as F
import torch
import torch.optim as optim

def check_image(path):
    try:
        im = Image.open(path)
        return True
    except:
        return False

transforms = transforms.Compose([
    transforms.Resize((64,64)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229,0.224,0.225])
])

train_data_path = './train'
train_data = torchvision.datasets.ImageFolder(root=train_data_path, transform=transforms)

val_data_path = './val'
val_data = torchvision.datasets.ImageFolder(root=val_data_path, transform=transforms)

test_data_path = './val'
test_data = torchvision.datasets.ImageFolder(root=test_data_path, transform=transforms)

batch_size = 64
train_data_loader = data.DataLoader(train_data, batch_size=batch_size)
val_data_loader = data.DataLoader(val_data, batch_size=batch_size)
test_data_loader = data.DataLoader(test_data, batch_size=batch_size)

device = 'cpu'

class CNNNet(nn.Module):

    def __init__(self, num_classes=2):
        super(CNNNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Linear(4096, num_classes)
        )
        self.flatten = torch.flatten

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = self.flatten(x, 1)
        x = self.classifier(x)
        return x



def train(model, optimizer, loss_fn, train_loader, val_loader, epochs=20, device="cpu"):
    for epoch in range(epochs):
        training_loss = 0.0
        valid_loss = 0.0
        model.train()
        for batch in train_loader:
            optimizer.zero_grad()
            inputs, targets = batch
            inputs = inputs.to(device)
            targets = targets.to(device)
            output = model(inputs)
            loss = loss_fn(output, targets)
            loss.backward()
            optimizer.step()
            training_loss += loss.data.item() * inputs.size(0)
        training_loss /= len(train_loader.dataset)

        model.eval()
        num_correct = 0
        num_examples = 0
        for batch in val_loader:
            inputs, targets = batch
            inputs = inputs.to(device)
            output = model(inputs)
            targets = targets.to(device)
            loss = loss_fn(output, targets)
            valid_loss += loss.data.item() * inputs.size(0)
            correct = torch.eq(torch.max(F.softmax(output), dim=1)[1], targets).view(-1)
            num_correct += torch.sum(correct).item()
            num_examples += correct.shape[0]
        valid_loss /= len(val_loader.dataset)

        print(
            'Epoch: , Training Loss: :.2f, Validation Loss: :.2f, accuracy = :.2f'.format(epoch, training_loss,
                                                                                                  valid_loss,
                                                                                                  num_correct / num_examples))

cnnnet = CNNNet()
optimizer = optim.Adam(cnnnet.parameters(), lr=0.001)

train(cnnnet, optimizer,torch.nn.CrossEntropyLoss(), train_data_loader,val_data_loader, epochs=10, device=device)

torch.save(cnnnet, './snet')

###### 下面是预测代码  #######

labels = ['cat','fish']

img = Image.open("./val/fish/100_1422.JPG") 
img = transforms(img).to(device)


prediction = F.softmax(cnnnet(img))
prediction = prediction.argmax()
print(labels[prediction]) 

这个例子还是很简单的