PyTorch快速入门教程六(使用LSTM做图片分类)

最新推荐文章于 2025-05-11 09:44:59 发布

转载

最新推荐文章于 2025-05-11 09:44:59 发布 · 2.9k 阅读

13 ·

CC 4.0 BY-SA版权

原文链接：https://my.oschina.net/earnp/blog/1113895

文章标签：

#人工智能 #python

本文介绍如何使用PyTorch中的LSTM进行图片分类。将28x28大小的图片视为长度为28的序列，通过LSTM模型进行处理。定义的模型包含2层LSTM和一个分类器，最终输出10类的概率。经过训练，模型在简单图像分类任务上表现出色。同时，提供了PyTorch的中文文档资源。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

为什么80%的码农都做不了架构师？>>>

对于LSTM，我们要处理的数据是一个序列数据，对于图片而言，我们如何将其转换成序列数据呢？图片的大小是28x28，所以我们可以将其看成长度为28的序列，序列中的每个数据的维度是28，这样我们就可以将其变成一个序列数据了。

model

class Rnn(nn.Module):
    def __init__(self, in_dim, hidden_dim, n_layer, n_class):
        super(Rnn, self).__init__()
        self.n_layer = n_layer
        self.hidden_dim = hidden_dim
        self.lstm = nn.LSTM(in_dim, hidden_dim, n_layer,
                            batch_first=True)
        self.classifier = nn.Linear(hidden_dim, n_class)

    def forward(self, x):
        # h0 = Variable(torch.zeros(self.n_layer, x.size(1),
                                #   self.h