批标准化BN层
导入必要的模块
import torch
import torch. nn as nn
import torch. nn. functional as F
import torch. optim as optim
import numpy as np
import matplotlib. pyplot as plt
from torchvision import transforms
"""
model: 针对一些常用的模型
datasets:用来下载一些经典的数据集
transforms: 提供了对PlL Image对象和Tensor对象的常用操作
utils:
"""
import torchvision
import os
导入数据
base_dir = r"./dataset/4weather"
train_dir = os. path. join( base_dir, "train" )
test_dir = os. path. join( base_dir, "test" )
配置训练集和测试集
transform = transforms. Compose( [
transforms. Resize( 96 , 96 ) ,
transforms. ToTensor( ) ,
transforms. Normalize( mean= [ 0.5 , 0.5 , 0.5 ] ,
std = [ 0.5 , 0.5 , 0.5 ] )
] )
train_ds = torchvision. datasets. ImageFolder(
train_dir,
transform = transform
)
test_ds = torchvision. datasets. ImageFolder(
test_dir,
transform = transform
)
BATCHSIZE = 16
train_dl = torch. utils. data. DataLoader(
train_ds,
batch_size = BATCHSIZE,
shuffle = True
)
test_dl = torch. utils. data. DataLoader(
test_dir,
batch_size = BATCHSIZE
)
创建网络层
class Net ( nn. Module) :
def __init__ ( self) :
super ( Net, self) . __init__( )
self. conv1 = nn. Conv2d( 3 , 16 , 3 )
self. bn1 = nn. BatchNorm2d( 16 )
self. pool = nn. MaxPool2d( 2 , 2 )
self. conv2 = nn. Conv2d( 16 , 32 , 2 )
self. bn2 = nn. BatchNorm2d( 32 )
self. conv3 = nn. Conv2d( 32 , 64 , 2 )
self. bn3 = nn. BatchNorm2d( 64 )
self. drop = nn. Dropout( 0.5 )
self. fc1 = nn. Linear( 64 * 10 * 10 , 1024 )
self. bn_fc1 = nn. BatchNorm2d( 1024 )
self. fc2 = nn. Linear( 1024 , 256 )
self. bn_fc2 = nn. BatchNorm2d( 256 )
self. fc3 = nn. Linear( 256 , 4 )
def forward ( self, x) :
x = self. pool( F. relu( self. conv1( x) ) )
x = self. bn1( x)
x = self. pool( F. relu( self. conv2( x) ) )
x = self. bn2( x)
x = self. pool( F. relu( self. conv3( x) ) )
x = x. view( - 1 , 64 * 10 * 10 )
x = F. relu( self. fc1( x) )
x = self. bn_fc1( x)
x = self. drop( x)
x = F. relu( self. fc2( x) )
x = self. bn_f2( x)
x = self. drop( x)
x = self. fc3( x)
return x
定义优化器和损失函数
model = Net( )
if torch. cuda. is_available( ) :
model. to( "cuda:0" )
optim = torch. optim. Adam( model. parameters( ) , lr = 0.001 )
loss_fn = nn. CrossEntropyLoss( )
定义评判标准
def fit ( epoch, model, trainloader, testloader) :
correct = 0
total = 0
running_loss = 0
model. train( )
for x, y in trainloader:
if torch. cuda. is_available( ) :
x, y = x. to( "cuda:0" ) , y. to( "cuda:0" )
y_pred = model( x)
loss = loss_fn( y_pred, y)
optim. zero_grad( )
loss. backward( )
optim. step( )
with torch. no_grad( ) :
y_pred = torch. argmax( y_pred, dim= 1 )
correct += ( y_pred == y) . sum ( ) . item( )
total += y. size( 0 )
running_loss += loss. item( )
epoch_loss = running_loss / len ( trainloader. dataset)
epoch_acc = correct / total
test_correct = 0
test_total = 0
test_running_loss = 0
model. eval ( )
with torch. no_grad( ) :
for x, y in testloader:
if torch. cuda. is_available( ) :
x , y = x. to( "cuda:0" ) , y. to( "cuda:0" )
y_pred = model( x)
test_correct += ( y_pred == y) . sum ( ) . item( )
test_total += y. size( 0 )
test_running_loss += loss. item( )
epoch_test_loss = test_running_loss / len ( testloader. dataset)
epoch_test_acc = test_correct / test_total
print ( "epoch:" , epoch,
"loss:" , round ( epoch_loss, 3 ) ,
"accuracy:" , round ( epoch_acc) ,
"test_loss:" , round ( epoch_test_loss, 3 ) ,
"test_accuracy:" , round ( epoch_test_acc, 3 )
)
return epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc
开启训练和测试
epochs = 30
train_loss = [ ]
train_acc = [ ]
test_loss = [ ]
test_acc = [ ]
for epoch in range ( epochs) :
epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc = fit( epoch,
model,
train_dl,
test_dl)
train_loss. append( epoch_loss)
train_acc. append( epoch_acc)
test_loss. append( epoch_test_loss)
test_acc. append( epoch_test_acc)