python学习day39

图像数据与显存

知识点回顾

1.图像数据的格式:灰度和彩色数据

2.模型的定义

3.显存占用的4种地方

a.模型参数+梯度参数

b.优化器参数

c.数据批量所占显存

d.神经元输出中间状态

4.batchisize和训练的关系

python 复制代码
import torch
import torchvision
import torch.nn as nn
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

torch.manual_seed(42)
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
#加载CIFAR10数据集
trainset = torchvision.datasets.CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform=transform
)
#创建数据加载器
train_loader = torch.utils.data.DataLoader(
    trainset,
    batch_size=4,
    download=True,
    shuffle=True
)
# CIFAR-10的10个类别
classes = ('plane', 'car', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

#随机图片
sample_idx = torch.randint(0, len(trainset), (1,)).item()
img, label = trainset[sample_idx]
#打印形状
print(img.shape)
print(classes[label])
#定义图像显示
def imshow(img):
    img = img / 2 + 0.5
    nping = img.numpy()
    plt.imshow(np.transpose(nping, (1, 2, 0)))
    plt.axis('off')
    plt.show()
imshow(img)


class MLP(nn.Module):
    def __init__(self, input_size=3072, hidden_size=128, output_size=10):
        super(MLP, self).__init__()
        self.flatten =  nn.Flatten()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)
    
    def forward(self, x):
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x
    
model = MLP()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

from torchsummary import summary
print("\n模型信息")
summary(model, (3, 32, 32))

OOM处理方案

显存占用部分

  1. 模型参数与梯度:模型的权重(Parameters)和对应的梯度(Gradients)会占用显存,尤其是深度神经网络(如 Transformer、ResNet 等),一个 1 亿参数的模型(如 BERT-base),单精度(float32)参数占用约 400MB(1e8×4Byte),加上梯度则翻倍至 800MB(每个权重参数都有其对应的梯度)。

  2. 部分优化器(如 Adam)会为每个参数存储动量(Momentum)和平方梯度(Square Gradient),进一步增加显存占用(通常为参数大小的 2-3 倍)

  3. 其他开销。

python 复制代码
#参数占用内存
"""
3.1模型参数与梯度参数
参数和梯度占用,二者大致相等
原来数据类型转化成float32 4B
"""
model = MLP()
total_params = sum(p.numel() for p in model.parameters())
print('Total parameters:', total_params)
print(f"Total parameters (float32): {total_params * 4 / 1024 / 1024:.2f}MB")
"""
3.2优化器参数
Adam优化器参数占用,存储有额外状态
"""

"""
3.3数据批量的显存占用
"""

"""
3.4前向/反向传播中间变量
"""
相关推荐
DKPT2 小时前
Java桥接模式实现方式与测试方法
java·笔记·学习·设计模式·桥接模式
子燕若水3 小时前
Unreal Engine 5中的AI知识
人工智能
极限实验室4 小时前
Coco AI 实战(一):Coco Server Linux 平台部署
人工智能
杨过过儿4 小时前
【学习笔记】4.1 什么是 LLM
人工智能
巴伦是只猫4 小时前
【机器学习笔记Ⅰ】13 正则化代价函数
人工智能·笔记·机器学习
伍哥的传说4 小时前
React 各颜色转换方法、颜色值换算工具HEX、RGB/RGBA、HSL/HSLA、HSV、CMYK
深度学习·神经网络·react.js
大千AI助手4 小时前
DTW模版匹配:弹性对齐的时间序列相似度度量算法
人工智能·算法·机器学习·数据挖掘·模版匹配·dtw模版匹配
AI生存日记4 小时前
百度文心大模型 4.5 系列全面开源 英特尔同步支持端侧部署
人工智能·百度·开源·open ai大模型
LCG元5 小时前
自动驾驶感知模块的多模态数据融合:时序同步与空间对齐的框架解析
人工智能·机器学习·自动驾驶