CIFAR10-Pytorch-94.1%

最新推荐文章于 2022-12-15 02:15:24 发布

Calcular

最新推荐文章于 2022-12-15 02:15:24 发布

阅读量703

点赞数

分类专栏：深度学习

本文链接：https://blog.csdn.net/Calcular/article/details/82384704

版权

该博客详细介绍了如何使用PyTorch训练模型，在CIFAR10数据集上达到了94.1%的测试准确率。内容包括`train.py`中的训练过程，`test.py`中的验证和测试步骤，以及`cifar10.py`中数据预处理和模型定义的代码实现。

摘要由CSDN通过智能技术生成

train.py:

# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import numpy as np
from torchvision import models
import torchvision.transforms as transforms
from cifar10 import CIFAR10

# define network
def genNet():
    net=models.resnet34(pretrained=True) # pretrained
    net.conv1=nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) #(b,64,32,32)
    net.bn1=nn.BatchNorm2d(64)
    net.maxpool=nn.MaxPool2d(kernel_size=1, stride=1)
    net.avgpool=nn.AvgPool2d(4, stride=1) #32/8
    expansion=1 #18,34 : expansion=1 else 4
    net.fc=nn.Linear(512 * expansion, 10) #
    return net

# define parameters
torch.cuda.set_device(0)
gpu_en=True
load=True
save=True
num_epochs=1
nbatch=150
lr=0.001

# define transform
transform=transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# define dataset
dataset=CIFAR10('D:\\dataset\\cifar10', train=True, transform=transform)

# run
net=genNet()
if load:
    net.load_state_dict(torch.load('net.pkl'))
if gpu_en:
    net=net.cuda()
loader=torch.utils.data.DataLoader(dataset, batch_size=nbatch, shuffle=True, num_workers=0)
criterion=nn.CrossEntropyLoss()
optimizer=torch.o