Ks dev

src/train.py

@@ -1,6 +1,3 @@
-# -*- coding: utf-8 -*-
-# 作者：小土堆
-# 公众号：土堆碎念
 
 import torchvision
 from torch.utils.tensorboard import SummaryWriter
@@ -10,6 +7,8 @@
 from torch import nn
 from torch.utils.data import DataLoader
 
+# 使用GPU
+device = torch.device("cuda")
 train_data = torchvision.datasets.CIFAR10(root="../data", train=True, transform=torchvision.transforms.ToTensor(),
                                           download=True)
 test_data = torchvision.datasets.CIFAR10(root="../data", train=False, transform=torchvision.transforms.ToTensor(),
@@ -29,9 +28,11 @@
 
 # 创建网络模型
 tudui = Tudui()
+tudui = tudui.to(device)
 
 # 损失函数
 loss_fn = nn.CrossEntropyLoss()
+loss_fn = loss_fn.to(device)
 
 # 优化器
 # learning_rate = 0.01
@@ -57,6 +58,8 @@
     tudui.train()
     for data in train_dataloader:
         imgs, targets = data
+        imgs = imgs.to(device)
+        targets = targets.to(device)
         outputs = tudui(imgs)
         loss = loss_fn(outputs, targets)
 
@@ -77,6 +80,8 @@
     with torch.no_grad():
         for data in test_dataloader:
             imgs, targets = data
+            imgs = imgs.to(device)
+            targets = targets.to(device)
             outputs = tudui(imgs)
             loss = loss_fn(outputs, targets)
             total_test_loss = total_test_loss + loss.item()