import torch
import torch.nn as nn
from vggnet import VGGNet
from loaddata import train_loader
import os
#防止核崩溃
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
epoch_num =200 #训练轮数
lr = 0.01 #学习率
batch_size =128#没批训练128个样本
net=VGGNet()
loss_func= nn.CrossEntropyLoss()
optimizer=torch.optim.Adam(net.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=5,gamma=0.9)
step_n=0
for epoch in range(epoch_num):
print(" epoch is",epoch)
for i,data in enumerate(train_loader):
inputs,labels=data
outputs = net(inputs) #前向传播
loss=loss_func(outputs,labels)
optimizer.zero_grad()#清空梯度
loss.backward() #反向传播,计算梯度
optimizer.step()#更新参数
_,pred=torch.max(outputs.data,dim=1)
correct=pred.eq(labels.data).cpu().sum()
print("step",i,';oss is',loss.item(),'mini-batch correct is:',
100.0*correct/batch_size)
if not os.path.exists('models'):
os.mkdir('models')
torch.save(net.state_dict(),'models/{}.pth'.format(epoch+1))
scheduler.step() #更新学习率
print("lr is,",optimizer.state_dict()["param_groups"][0]["lr"])
这里的代码能用.cuda()加速吗?应该怎么改呀老师