pytorch tips：使pytorch代码在不改动情况在有GPU自动在GPU运行

最新推荐文章于 2025-03-12 21:49:47 发布

小贤2016

最新推荐文章于 2025-03-12 21:49:47 发布

阅读量7.5k

点赞数 4

分类专栏：计算机视觉 python pytorch 文章标签： python pytorch resnet 计算机视觉

本文链接：https://blog.youkuaiyun.com/u013362955/article/details/83189222

版权

定义：

#gpu or not
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

使用：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import argparse
import torch.utils.data
from resnet import ResNet18

#gpu or not
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

parser = argparse.ArgumentParser(description='Python CIFAR10 Training')
parser.add_argument('--outf', default='./model', help='folder to output images and model checkpoints')
parser.add_argument('--net', default=',.model/Resnet18.pth', help='path to net(to continue training)')
args = parser.parse_args()

#超参数设置
EPOCH = 135
pre_epoch = 0
BATCH_SIZE = 128
LR