Pytorch调用GPU实现简单的卷积神经网络

"""
一个简单的调用GPU的卷积神经网络示例!
"""

import torch
import torchvision
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader

batch_size = 64
device = torch.device("cuda:0")

class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.convl1 = torch.nn.Conv2d(1,10,5)
        self.convl2 = torch.nn.Conv2d(10,20,5)
        self.pooling = torch.nn.MaxPool2d(2)
        self.activate = torch.nn.ReLU()
        self.linear = torch.nn.Linear(320,10)

    def forward(self,x):
        x = x.view(-1,1,28,28)
        x = self.convl1(x)
        x = self.pooling(x)
        x = self.activate(x)
        x = self.convl2(x)
        x = self.pooling(x)
        x = x.view(-1,320)
        x = self.linear(x)
        return x

def train(train_loader, model, criterion, optimizer, epoch):
    loss_sum = 0.0
    for index, (x, y) in enumerate(train_loader):
        x = x.to(device)
        y = y.to(device)
        y_hat = model(x)
        loss = criterion(y_hat, y)
        loss_sum += loss.item()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (index % 300 == 299):
            print(epoch,",",index//300, ":", loss_sum/300)
            loss_sum = 0.0


def test(test_loader, model):
    total = 0
    correct = 0
    for x,y in test_loader:
        x = x.to(device)
        y = y.to(device)
        y_hat = model(x)
        _,guess = torch.max(y_hat,1)
        correct += (guess == y).sum().item()
        total += y.size(0)
    print("ACC == ", correct / total)


if __name__ == '__main__':
    transformer=torchvision.transforms.Compose([torchvision.transforms.ToTensor(),torchvision.transforms.Normalize(0.1307, 0.3018)])
    train_data = datasets.MNIST('MNIST',True,transformer,download=True)
    test_data=datasets.MNIST('MNIST',True,transformer,download=True)

    train_loader = DataLoader(train_data,batch_size=batch_size,shuffle=True, num_workers=2)
    test_loader = DataLoader(test_data,batch_size=batch_size,shuffle=True,num_workers=2)
    model = CNN()
    model.to(device)
    criterion = torch.nn.CrossEntropyLoss()
    criterion.to(device)
    optimizer = torch.optim.SGD(model.parameters(),lr=0.01,momentum=0.5)
    for epoch in range(10):
        train(train_loader,model,criterion,optimizer,epoch)

    test(test_loader,model)

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

DeniuHe

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值