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

最新推荐文章于 2025-06-26 11:17:17 发布

小贤2016

最新推荐文章于 2025-06-26 11:17:17 发布

阅读量7.5k

点赞数 4

CC 4.0 BY-SA版权

分类专栏：计算机视觉 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 = 0.1

#准备数据集并预处理
transform_train = transforms.Compose([
    transforms.RandomCrop(size=32, padding=4), #先padding，再随机截取32*32
    transforms.RandomHor