文章目录
从头搭建一个深度学习的模型,基本上都可以从这个框架去套用。
包括了最基础的模型的定义、训练、验证和测试
完整的框架代码
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader,Dataset
# 定义模型
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.fc1=nn.Linear(100,50)
self.fc2=nn.Linear(50,10)
def forward(self,x):
x=torch.relu(self.fc1(x))
x=self.fc2(x)
return x
# 定义训练集
class TrainDataset(Dataset):
def __init__(self,data_path,target_path):
super(TrainDataset,self).__init__()
pass # 初始化
def __getitem__(self, index):
pass# 根据index返回数据
# return data,target
def __len__(self):
pass # 返回数据的长度
# return len(xxx)
# 定义验证集
class ValDataset(Dataset):
def __init__(self, data_path, target_path):
super(ValDataset, self).__init__()
pass # 初始化
def __getitem__(self, index):
pass # 根据index返回数据
# return data,target
def __len__(self):
pass # 返回数据的长度
# return len(xxx)
# 定义测试集
class TestDataset(Dataset):
def __init__(self, data_path, target_path):
super(TestDataset, self).__init__()
pass # 初始化
def __getitem__(self, index):
pass # 根据index返回数据
# return data,target
def __len__(self):
pass # 返回数据的长度
# return len(xxx)
# 训练
def train(model,optimizer,criterion,train_loader,device):
total_loss=0
# 将模型设置为训练模式
model.train()
# 遍历数据集
for data,target in train_loader:
# 将数据移到设备上
data,target=data.to(device),target.to(device)
# 梯度清零
optimizer.zero_grad()
# 前向传播
output=model(data)
# 计算损失
loss=criterion(output,target)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
# 统计损失
total_loss+=loss.item()
# 计算平均损失
avg_loss=total_loss/len(train_loader)
return avg_loss
# 验证
def validate(model,criterion,val_loader,device):
total_loss=0
# 将模型设置为验证模式
model.eval()
with torch.no_grad():
#遍历验证集
for data,target in val_loader:
# 将数据移到设备上
data,target=data.to(device),target.to(device)
# 前向传播
output=model(data)
# 计算损失
loss=criterion(output,target)
# 统计损失
total_loss+=loss.item()
# 计算平均损失
avg_loss=total_loss/len(val_loader)
return avg_loss
def test(model,criterion,test_loader,device):
total_loss=0
total_correct=0
# 将模型设置为测试模式
model.eval()
with torch.no_grad():
# 遍历测试集
for data,target in test_loader:
# 将数据迁移到设备上
data,target=data.to(device),target.to(device)
# 前向传播
output=model(data)
# 计算损失
loss=criterion(output,target)
# 统计损失
total_loss+=loss.item()
# 计算准确率
_,predicted =torch.max(output.data,1)
total_correct+=(predicted==target).sum().item()
# 计算平均损失
avg_loss=total_loss/len(test_loader)
# 计算准确率
accuracy=total_correct/len(test_loader.dataset)
return avg_loss,accuracy
# 主函数
def main():
# 定义超参数
epoch=10
lr=0.01
batch_size=32
# 定义设备
device =torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 定义训练集、验证集、测试集 的数据
train_data=TrainDataset(data_path='',target_path='')
val_data = ValDataset(data_path='', target_path='')
test_data = TestDataset(data_path='', target_path='')
# 定义训练集、验证集、测试集 的数据加载器
train_loader=DataLoader(train_data,batch_size=batch_size,shuffle=True)
val_loader=DataLoader(val_data,batch_size=batch_size,shuffle=True)
test_loader=DataLoader(test_data,batch_size=batch_size,shuffle=True)
# 定义模型
model=Model().to(device)
# 加载预训练权重
model.load_state_dict(torch.load('xxx.pth')
# 定义损失函数和优化器
criterion=nn.CrossEntropyLoss()
optimizer=optim.SGD(model.parameters(),lr=lr)
# 训练和验证
for epoch in range(1,epoch+1):
train_loss=train(model,optimizer,criterion,train_loader,device)
val_loss=validate(model,criterion,val_loader,device)
# 显示训练集和验证集的损失
if (epoch%display_iter)==0:
print(f"Epoch: {epoch}, Train loss:{train_loss:.4f},Val loss: {val_loss: .4f}")
# 保存权重
if (epoch%snapshot_iter)==0:
torch.save(model.state_dict(),"xxx.pth")
# 测试
test_loss,test_accuracy=test(model,criterion,test_loader,device)
print(f"Test loss:{test_loss:.4f},Test accuracy:{test_accuracy:.4f}")
if __name__=='__main__':
main()