Pytorch实战 | P2 彩色图片识别(深度学习实践pytorch)

陈大麦 .

已于 2022-10-18 13:24:40 修改

阅读量945

点赞数 1

分类专栏：深度学习实践100例文章标签：深度学习 pytorch python

于 2022-10-18 11:57:26 首次发布

本文链接：https://blog.csdn.net/weixin_35820421/article/details/127383945

版权

一、我的环境：

● 语言环境：Python3.8
● 编译器：pycharm
● 深度学习环境：Pytorch

二、主要代码实现

1、main.py

import matplotlib.pyplot as plt
import torch
import numpy as np
from model import *
from torchinfo import summary

# 一、 数据准备

# --- 1、设置GPU ---
import torchvision.datasets

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --- 2、导入数据 ---
train_ds = torchvision.datasets.CIFAR10('data',
                                        train=True,
                                        transform=torchvision.transforms.ToTensor(),  # 将数据类型转化为Tensor
                                        download=True)

test_ds = torchvision.datasets.CIFAR10('data',
                                       train=False,
                                       transform=torchvision.transforms.ToTensor(),  # 将数据类型转化为Tensor
                                       download=False)

batch_size = 32

train_dl = torch.utils.data.DataLoader(train_ds,
                                       batch_size=batch_size,
                                       shuffle=False)
test_dl = torch.utils.data.DataLoader(test_ds, batch_size=batch_size)

# --- 3、数据可视化 ---
imgs, labels = next(iter(train_dl))