Spatial Transformer Network实现

最新推荐文章于 2024-05-15 10:09:05 发布

转载最新推荐文章于 2024-05-15 10:09:05 发布 · 319 阅读

0 ·

CC 4.0 BY-SA版权

原文链接：https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html

文章标签：

#transformer #pytorch #深度学习

该博客介绍了如何使用PyTorch实现空间变换网络（Spatial Transformer Network）。网络包括局部定位网络和仿射矩阵回归器，用于在卷积神经网络中进行空间变换。代码展示了一个完整的前向传播过程，包括输入的变换和常规的卷积操作。

部署运行你感兴趣的模型镜像

import torch
import torch.nn as nn
import torch.nn.functional as F


class SpatialTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

        # Spatial transformer localization-network
        self.localization = nn.Sequential(
            nn.Conv2d(1, 8, kernel_size=7),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
            nn.Conv2d(8, 10, kernel_size=5),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
        )

        # Regressor for the 3 * 2 affine matrix
        self.fc_loc = nn.Sequential(
            nn.Linear(10 * 3 * 3, 32), nn.ReLU(True), nn.Linear(32, 3 * 2)
        )

        # Initialize the weights/bias with identity transformation
        self.fc_loc[2].weight.data.zero_()
        self.fc_loc[2].bias.data.copy_(
            torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float)
        )

    # Spatial transformer network forward function
    def stn(self, x):
        xs = self.localization(x)
        xs = xs.view(-1, 10 * 3 * 3)
        theta = self.fc_loc(xs)
        theta = theta.view(-1, 2, 3)

        grid = F.affine_grid(theta, x.size())
        x = F.grid_sample(x, grid)

        return x

    def forward(self, x):
        # transform the input
        x = self.stn(x)

        # Perform the usual forward pass
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)


def primal_spatial_transformer():
    # Reference:
    # [1] https://proceedings.neurips.cc//paper/2015/file/33ceb07bf4eeb3da587e268d663aba1a-Paper.pdf
    # [2] https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html
    return SpatialTransformer()


if __name__ == "__main__":
    model = primal_spatial_transformer()
    y = model.forward(torch.rand(4, 1, 28, 28))
    print(y.shape)

pytorch官网上有哈

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch

Cuda

PyTorch 是一个开源的 Python 机器学习库，基于 Torch 库，底层由 C++ 实现，应用于人工智能领域，如计算机视觉和自然语言处理