import numpy as np
import torch
import torch. nn as nn
import torch. optim as optim
import torchvision
import torchvision. transforms as transforms
import torch. nn. functional as F
import torch. optim as optim
from torchvision import datasets, transforms
import os
import time
print ( "PyTorch Version: " , torch. __version__)
图像增强
data_transform = transforms. Compose( [
transforms. Resize( 40 ) ,
transforms. RandomHorizontalFlip( ) ,
transforms. RandomCrop( 32 ) ,
transforms. ToTensor( )
] )
trainset = torchvision. datasets. ImageFolder( root= '/home/kesci/input/CIFAR102891/cifar-10/train'
, transform= data_transform)
trainset[ 0 ] [ 0 ] . shape
data = [ d[ 0 ] . data. cpu( ) . numpy( ) for d in trainset]
np. mean( data)
np. std( data)
transform_train = transforms. Compose( [
transforms. RandomCrop( 32 , padding= 4 ) ,
transforms. RandomHorizontalFlip( ) ,
transforms. ToTensor( ) ,
transforms. Normalize( ( 0.4731 , 0.4822 , 0.4465 ) , ( 0.2212 , 0.1994 , 0.2010 ) ) ,
] )
transform_test = transforms. Compose( [
transforms. ToTensor( ) ,
transforms. Normalize( ( 0.4731 , 0.4822 , 0.4465 ) , ( 0.2212 , 0.1994 , 0.2010 ) ) ,
] )
导入数据集
train_dir = '/home/kesci/input/CIFAR102891/cifar-10/train'
test_dir = '/home/kesci/input/CIFAR102891/cifar-10/test'
trainset = torchvision. datasets. ImageFolder( root= train_dir, transform= transform_train)
trainloader = torch. utils. data. DataLoader( trainset, batch_size= 256 , shuffle= True )
testset = torchvision. datasets. ImageFolder( root= test_dir, transform= transform_test)
testloader = torch. utils. data. DataLoader( testset, batch_size= 256 , shuffle= False )
classes = [ 'airplane' , 'automobile' , 'bird' , 'cat' , 'deer' , 'dog' , 'forg' , 'horse' , 'ship' , 'truck' ]
class ResidualBlock ( nn. Module) :
def __init__ ( self, inchannel, outchannel, stride= 1 ) :
super ( ResidualBlock, self) . __init__( )
self. left = nn. Sequential(
nn. Conv2d( inchannel, outchannel, kernel_size= 3 , stride= stride, padding= 1 , bias= False ) ,
nn. BatchNorm2d( outchannel) ,
nn. ReLU( inplace= True ) ,
nn. Conv2d( outchannel, outchannel, kernel_size= 3 , stride= 1 , padding= 1 , bias= False ) ,
nn. BatchNorm2d( outchannel)
)
self. shortcut = nn. Sequential( )
if stride != 1 or inchannel != outchannel:
self. shortcut = nn. Sequential(
nn. Conv2d( inchannel, outchannel, kernel_size= 1 , stride= stride, bias= False ) ,
nn. BatchNorm2d( outchannel)
)
def forward ( self, x) :
out = self. left( x)
out += self. shortcut( x)
out = F. relu( out)
return out
class ResNet ( nn. Module) :
def __init__ ( self, ResidualBlock, num_classes= 10 ) :
super ( ResNet, self) . __init__( )
self. inchannel = 64
self. conv1 = nn. Sequential(
nn. Conv2d( 3 , 64 , kernel_size= 3 , stride= 1 , padding= 1 , bias= False ) ,
nn. BatchNorm2d( 64 ) ,
nn. ReLU( ) ,
)
self. layer1 = self. make_layer( ResidualBlock, 64 , 2 , stride= 1 )
self. layer2 = self. make_layer( ResidualBlock, 128 , 2 , stride= 2 )
self. layer3 = self. make_layer( ResidualBlock, 256 , 2 , stride= 2 )
self. layer4 = self. make_layer( ResidualBlock, 512 , 2 , stride= 2 )
self. fc = nn. Linear( 512 , num_classes)
def make_layer ( self, block, channels, num_blocks, stride) :
strides = [ stride] + [ 1 ] * ( num_blocks - 1 )
layers = [ ]
for stride in strides:
layers. append( block( self. inchannel, channels, stride) )
self. inchannel = channels
return nn. Sequential( * layers)
def forward ( self, x) :
out = self. conv1( x)
out = self. layer1( out)
out = self. layer2( out)
out = self. layer3( out)
out = self. layer4( out)
out = F. avg_pool2d( out, 4 )
out = out. view( out. size( 0 ) , - 1 )
out = self. fc( out)
return out
def ResNet18 ( ) :
return ResNet( ResidualBlock)
训练和测试
device = torch. device( "cuda" if torch. cuda. is_available( ) else "cpu" )
EPOCH = 20
pre_epoch = 0
LR = 0.1
net = ResNet18( ) . to( device)
criterion = nn. CrossEntropyLoss( )
optimizer = optim. SGD( net. parameters( ) , lr= LR, momentum= 0.9 , weight_decay= 5e - 4 )
if __name__ == "__main__" :
print ( "Start Training, Resnet-18!" )
num_iters = 0
for epoch in range ( pre_epoch, EPOCH) :
print ( '\nEpoch: %d' % ( epoch + 1 ) )
net. train( )
sum_loss = 0.0
correct = 0.0
total = 0
for i, data in enumerate ( trainloader, 0 ) :
num_iters += 1
inputs, labels = data
inputs, labels = inputs. to( device) , labels. to( device)
optimizer. zero_grad( )
outputs = net( inputs)
loss = criterion( outputs, labels)
loss. backward( )
optimizer. step( )
sum_loss += loss. item( ) * labels. size( 0 )
_, predicted = torch. max ( outputs, 1 )
total += labels. size( 0 )
correct += ( predicted == labels) . sum ( ) . item( )
if ( i + 1 ) % 20 == 0 :
print ( '[epoch:%d, iter:%d] Loss: %.03f | Acc: %.3f%% '
% ( epoch + 1 , num_iters, sum_loss / ( i + 1 ) , 100 . * correct / total) )
print ( "Training Finished, TotalEPOCH=%d" % EPOCH)