模型权重保存
导入必要模型
import torch
import torch. nn as nn
import torch. nn. functional as F
import torch. optim as optim
import matplotlib. pyplot as plt
import numpy as np
import torchvision
from torchvision import transforms
from torch. utils. data import DataLoader
import os
import copy
导入数据
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= [ .5 , .5 , .5 ] ,
std= [ .5 , .5 , .5 ]
)
] )
train_ds = torchvision. datasets. ImageFolder(
train_dir,
transform= transform
)
test_ds = torchvision. datasets. ImageFolder(
test_dir,
transform= transform
)
BATCH_SIZE = 32
train_dl = DataLoader(
train_ds,
batch_size= BATCH_SIZE,
shuffle= True
)
test_dl = DataLoader(
test_ds,
batch_size= BATCH_SIZE,
)
定义网络模型类
class Net ( nn. Module) :
def __init__ ( self) :
super ( Net, self) . __init__( )
self. conv1 = nn. Conv2d( 3 , 16 , 3 )
self. pool = nn. MaxPool2d( 2 , 2 )
self. conv2 = nn. Conv2d( 16 , 32 , 3 )
self. conv3 = nn. Conv2d( 32 , 64 , 3 )
self. fc1 = nn. Linear( 64 * 10 * 10 , 1024 )
self. fc2 = nn. Linear( 1024 , 256 )
self. fc3 = nn. Linear( 256 , 4 )
def forward ( self, x) :
x = self. pool( F. relu( self. conv1( x) ) )
x = self. pool( F. relu( self. conv2( x) ) )
x = self. pool( F. relu( self. conv3( x) ) )
x = x. view( - 1 , 64 * 10 * 10 )
x = F. relu( self. fc1( x) )
x = F. relu( self. fc2( x) )
x = self. fc3( x)
return x
定义损失函数和优化器
model = Net( )
if torch. cuda. is_available( ) :
model. to( "cuda" )
optimizer = 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" ) , y. to( "cuda" )
y_pred = model( x)
loss = loss_fn( y_pred, y)
optimizer. zero_grad( )
loss. backward( )
optimizer. 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" ) , y. to( "cuda" )
y_pred = model( x)
loss = loss_fn( y_pred, y)
y_pred = torch. argmax( y_pred, dim= 1 )
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, 3 ) ,
"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 = 100
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)
plt. plot( range ( 1 , epochs + 1 ) , train_loss, label= 'train_loss' )
plt. plot( range ( 1 , epochs + 1 ) , test_loss, label= 'test_loss' )
plt. legend( )
模型保存
"""
保存模型
staic_dict 就是一个简单的Python字典,它将模型中的可训练参数(比如:weights和biases, batchnorm的running_mean、
torch.optim参数等)通过将模型每层与层的参数张量之间一一映射,实现保存、更新、变化和再存储。
"""
PATH = "./my_net.pth"
torch. save( model. state_dict( ) , PATH)
加载模型
new_model = Net( )
new_model. load_state_dict( torch. load( PATH) )
new_model. to( "cuda" )
测试
test_correct = 0
test_total = 0
new_model. eval ( )
with torch. no_grad( ) :
for x, y in test_dl:
if torch. cuda. is_available( ) :
x, y = x. to( "cuda" ) , y. to( "cuda" )
y_pred = new_model( x)
y_pred = torch. argmax( y_pred, dim= 1 )
test_correct += ( y_pred== y) . sum ( ) . item( )
test_total += y. size( 0 )
epoch_test_acc = test_correct / test_total
print ( epoch_test_acc)
训练函数保存最优参数
model = Net( )
if torch. cuda. is_available( ) :
model. to( "cuda" )
optim = torch. optim. Adam( model. parameters( ) , lr= 0.001 )
best_model_wts = copy. deepcopy( model. state_dict( ) )
best_acc = 0.0
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)
if epoch_test_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy. deepcopy( model. state_dict( ) )
model. load_state_dict( best_model_wts)
model. eval ( )
完整模型的保存和加载
PATH = r"./my_whole_model.pth"
torch. save( model, PATH)
new_model2 = torch. load( PATH)
new_model2. eval ( )
PATH = "./my_gpu_model_wts"
torch. save( model. state_dict( ) , PATH)
device = torch. device( "cpu" )
model = Net( )
model. load_state_dict( torch. load( PATH, map_location= device) )
PATH = r"./my_gpu_model2_wts"
torch. save( model. state_dict( ) , PATH)
device = Net( )
model. load_state_dict( torch. load( PATH) )
model. to( device)