import torch
import numpy as np

# 直接构造
data = [[1, 2],[3, 4]]
x_data = torch.tensor(data)

# 使用numpy array转换
np_array = np.array(data)
x_np = torch.from_numpy(np_array)

# 构造一个指定大小的Tensors
shape = (2,3,)
rand_tensor = torch.rand(shape)		# 生成一个2x3的Tensors，内容是随机
ones_tensor = torch.ones(shape)		# 生成一个2x3的Tensors，内容全为1
zeros_tensor = torch.zeros(shape)	# 生成一个2x3的Tensors，内容全为0

# 生成从0到9的tensor
a = torch.arange(10)	# tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

# 矩阵乘法
C = torch.matmul(A, B)	

# 沿着纬度 dim 进行拼接
tensor1 = torch.tensor([[1, 2, 3, 4],
                        [5, 6, 7, 8],
                        [9, 10, 11, 12]])

tensor2 = torch.tensor([[13, 14, 15, 16],
                        [17, 18, 19, 20],
                        [21, 22, 23, 24]])
                        
result1 = torch.cat([tensor1, tensor2], dim=0)
# tensor([[ 1,  2,  3,  4],
#         [ 5,  6,  7,  8],
#         [ 9, 10, 11, 12],
#         [13, 14, 15, 16],
#         [17, 18, 19, 20],
#         [21, 22, 23, 24]])

result2 = torch.cat([tensor1, tensor2], dim=1)
# tensor([[ 1,  2,  3,  4, 13, 14, 15, 16],
#         [ 5,  6,  7,  8, 17, 18, 19, 20],
#         [ 9, 10, 11, 12, 21, 22, 23, 24]])

a = torch.tensor([1, 2, 3])

# 重复每个元素两次
b = torch.repeat_interleave(a, 2)		# tensor([1, 1, 2, 2, 3, 3])

# 根据tensor重复
repeats = torch.tensor([2, 3, 1])
c = torch.repeat_interleave(a, repeats)		# tensor([1, 1, 2, 2, 2, 3])

torch.save(model, 'model.pth')
model = torch.load('model.pth')

# 继承 nn.Module，实现一个Attention
class Attention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
        
    def forward(self, query, key, value):
        # Implement attention mechanism
        pass

# 将多层网络按顺序级联
model = nn.Sequential(
    nn.Conv2d(1, 20, 5),
    nn.ReLU(),
    nn.Conv2d(20, 64, 5),
    nn.ReLU()
)
# 等同于
model2 = nn.Sequential()
model2.add_module("conv1", nn.Conv2d(1, 20, 5))
model2.add_module('relu1', nn.ReLU())
model2.add_module('conv2', nn.Conv2d(20, 64, 5))
model2.add_module('relu2', nn.ReLU())

class Encoder(nn.Module):
  def __init__(self):
    self.down_blocks = nn.ModuleList([])
    for i in xxx:
      down_block = ResnetDownsampleBlock2D(...)
      self.down_blocks.append(down_block)
  # nn.ModuleList需要自定处理前向传播，这里循环处理了sample，自由度更高，更适合复杂的库作者
  def forward(self, sample: torch.Tensor) -> torch.Tensor:
    def create_custom_forward(module):
        def custom_forward(*inputs):
            return module(*inputs)
        return custom_forward
    for down_block in self.down_blocks:
        sample = torch.utils.checkpoint.checkpoint(
            create_custom_forward(down_block), sample, use_reentrant=False
        ) 
    return sample

# 定义一个全连接层（线性层），将输入纬度映射到输出纬度
query = nn.Linear(in_features, out_features)

position_encoding = nn.Embedding(targets_length, d_model)

# 定义一个二维卷积层
conv = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1)

# 定义一个逆卷积层，用于放大特征图，实现上采样
conv_transpose = nn.ConvTranspose2d(in_channels=1, out_channels=1, kernel_size=3, stride=2, padding=1)

dropout = nn.Dropout(p=dropout_rate)