学习记录
import torch
from model_mnist.Net import ConvNet
from data.data_ import transform
from torch.utils.data import DataLoader
from torchvision.datasets import mnist
from torchvision import transforms
import time
MOMENTUM=0.5
LR = 0.01
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize([0.1307],[0.3081])
# 对于每个通道,从该通道的每个像素值中减去均值`0.1307`。
#然后,将结果除以该通道的标准差`0.3081`。
]
)
if torch.cuda.is_available():
print(f"Using GPU: {torch.cuda.get_device_name(0)}") # 0是设备索引,通常是主设备
else:
print("Using CPU")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_set = mnist.MNIST('./data', train=True, transform=transform, download=False)
train_loader = DataLoader(dataset = train_set,shuffle=True,batch_size=128,drop_last=True)
test_set = mnist.MNIST('./data', train=False, transform=transform, download=False)
test_loader = DataLoader(dataset = test_set,shuffle=True,batch_size=64,drop_last=True)
#实例化模型
model = ConvNet()
model.to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),lr = LR, momentum=MOMENTUM)
def train(epoch):
'''训练'''
for index, data in enumerate(train_loader, 0):
inputs, target = data# 取数据
inputs = inputs.to(device) # 将输入数据放到GPU上
#print(inputs.device)
target = target.to(device) # 将标签放到GPU上
#print(target.device)
optimizer.zero_grad() # 梯度清零
outputs = model(inputs) # 训练
loss = criterion(outputs, target) # 算损失
loss.backward() # 反向传播
optimizer.step() # 优化梯度
# 累计loss ,这里要用item()取数据要不回构建计算图
if epoch% 20== 19:
torch.save(model.state_dict(), "model_mnist/model_GPU.pth") # 保存模型
torch.save(optimizer.state_dict(), "model_mnist/optimize_GPUr.pth")
print("训练次数为:{},损失值为:{}".format(index, loss.item()))
if __name__ == '__main__':
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
start = time.time()
for epoch in range(20):
print({"————————第{}轮测试开始——————".format(epoch+ 1)})
train(epoch)
end = time.time()
print('Time cost:', end - start, 's')
batch_size =64 训练20轮 CPU与GPU对比
bacth_size=128 训练20轮 CPU与GPU对比
GPU缩短了一半时间