Pytorch学习14_线性层及其他层-优快云博客

本文链接：https://blog.youkuaiyun.com/weixin_53384447/article/details/135900747

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader

dataset=torchvision.datasets.CIFAR10("./dataset",train=False,
                                     transform=torchvision.transforms.ToTensor(),
                                     download=True)

dataloader=DataLoader(dataset,batch_size=64)

for data in dataloader:
    imgs,target=data
    print(f"imgs.shape:{imgs.shape}")

    output=torch.reshape(imgs,(1,1,1,-1))
    print(f"output.shape:{output.shape}")

使用 torch.reshape 进行形状重塑时，如果某一个维度使用了 -1，它表示该维度的大小将由系统自动计算，以确保总的元素数量不变。在 (1, 1, 1, -1) 中：

第一个维度：1
第二个维度：1
第三个维度：1
第四个维度：-1

这里的 -1 表示这个维度的大小将由系统自动计算。系统将会根据原始张量的元素总数，以及前面三个维度的大小，来计算这个维度的大小，以保持总的元素数量不变。

这种用法通常用于将多维张量重塑为特定形状，同时保持元素总数不变。

import torch
import torchvision
from torch import nn
from torch.nn import Linear
from torch.utils.data import DataLoader

# 加载 CIFAR-10 数据集，设置转换为张量并下载数据
dataset = torchvision.datasets.CIFAR10("./dataset", train=False,
                                       transform=torchvision.transforms.ToTensor(),
                                       download=True)

# 创建数据加载器，每批次包含64张图像
dataloader = DataLoader(dataset, batch_size=64)

# 定义一个神经网络类 nn_lin，继承自 nn.Module
class nn_lin(nn.Module):

    def __init__(self) -> None:
        super().__init__()
        # 定义一个线性层，输入大小为 196608，输出大小为 10
        self.linear01 = Linear(196608, 10)

    def forward(self, input):
        # 在前向传播中应用线性层
        output = self.linear01(input)
        return output

# 创建 nn_lin 类的实例
lin_xx = nn_lin()
step = 0

# 遍历数据加载器中的第一个批次
for data in dataloader:
    # 获取图像和目标标签
    imgs, target = data
    print(f"imgs.shape: {imgs.shape}")

    # 使用 torch.reshape 将图像展平成 (1, 1, 1, -1) 的形状
    output = torch.reshape(imgs, (1, 1, 1, -1))
    print(f"output.shape: {output.shape}")

    # 将展平后的图像传递给神经网络进行前向传播
    output_lin = lin_xx(output)
    print(f"output_lin: {output_lin}")
    
    # 仅遍历第一个批次以防止输出过多
    break