meissonic_temp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformer2DModel(\n",
      "  (pos_embed): FluxPosEmbed()\n",
      "  (time_text_embed): CombinedTimestepTextProjEmbeddings(\n",
      "    (time_proj): Timesteps()\n",
      "    (timestep_embedder): TimestepEmbedding(\n",
      "      (linear_1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "      (act): SiLU()\n",
      "      (linear_2): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "    )\n",
      "    (text_embedder): PixArtAlphaTextProjection(\n",
      "      (linear_1): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "      (act_1): SiLU()\n",
      "      (linear_2): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "    )\n",
      "  )\n",
      "  (context_embedder): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "  (transformer_blocks): ModuleList(\n",
      "    (0-13): 14 x TransformerBlock(\n",
      "      (norm1): AdaLayerNormZero(\n",
      "        (silu): SiLU()\n",
      "        (linear): Linear(in_features=1024, out_features=6144, bias=True)\n",
      "        (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=False)\n",
      "      )\n",
      "      (norm1_context): AdaLayerNormZero(\n",
      "        (silu): SiLU()\n",
      "        (linear): Linear(in_features=1024, out_features=6144, bias=True)\n",
      "        (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=False)\n",
      "      )\n",
      "      (attn): Attention(\n",
      "        (norm_q): RMSNorm()\n",
      "        (norm_k): RMSNorm()\n",
      "        (to_q): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "        (to_k): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "        (to_v): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "        (add_k_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "        (add_v_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "        (add_q_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "        (to_out): ModuleList(\n",
      "          (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "          (1): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "        (to_add_out): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "        (norm_added_q): RMSNorm()\n",
      "        (norm_added_k): RMSNorm()\n",
      "      )\n",
      "      (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=False)\n",
      "      (ff): FeedForward(\n",
      "        (net): ModuleList(\n",
      "          (0): GELU(\n",
      "            (proj): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "          )\n",
      "          (1): Dropout(p=0.0, inplace=False)\n",
      "          (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "        )\n",
      "      )\n",
      "      (norm2_context): LayerNorm((1024,), eps=1e-06, elementwise_affine=False)\n",
      "      (ff_context): FeedForward(\n",
      "        (net): ModuleList(\n",
      "          (0): GELU(\n",
      "            (proj): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "          )\n",
      "          (1): Dropout(p=0.0, inplace=False)\n",
      "          (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (single_transformer_blocks): ModuleList(\n",
      "    (0-27): 28 x SingleTransformerBlock(\n",
      "      (norm): AdaLayerNormZeroSingle(\n",
      "        (silu): SiLU()\n",
      "        (linear): Linear(in_features=1024, out_features=3072, bias=True)\n",
      "        (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=False)\n",
      "      )\n",
      "      (proj_mlp): Linear(in_features=1024, out_features=4096, bias=True)\n",
      "      (act_mlp): GELU(approximate='tanh')\n",
      "      (proj_out): Linear(in_features=5120, out_features=1024, bias=True)\n",
      "      (attn): Attention(\n",
      "        (norm_q): RMSNorm()\n",
      "        (norm_k): RMSNorm()\n",
      "        (to_q): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "        (to_k): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "        (to_v): Linear(in_features=1024, out_features=1024, bias=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (embed): UVit2DConvEmbed(\n",
      "    (embeddings): Embedding(8256, 1024)\n",
      "    (layer_norm): RMSNorm()\n",
      "    (conv): Conv2d(1024, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "  )\n",
      "  (mlm_layer): ConvMlmLayer(\n",
      "    (conv1): Conv2d(1024, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "    (layer_norm): RMSNorm()\n",
      "    (conv2): Conv2d(1024, 8192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "  )\n",
      "  (cond_embed): TimestepEmbedding(\n",
      "    (linear_1): Linear(in_features=2304, out_features=1024, bias=False)\n",
      "    (act): SiLU()\n",
      "    (linear_2): Linear(in_features=1024, out_features=1024, bias=False)\n",
      "  )\n",
      "  (encoder_proj_layer_norm): RMSNorm()\n",
      "  (project_to_hidden_norm): RMSNorm()\n",
      "  (project_to_hidden): Linear(in_features=1024, out_features=1024, bias=False)\n",
      "  (project_from_hidden_norm): RMSNorm()\n",
      "  (project_from_hidden): Linear(in_features=1024, out_features=1024, bias=False)\n",
      "  (down_block): Simple_UVitBlock(\n",
      "    (downsample): Downsample2D(\n",
      "      (norm): RMSNorm()\n",
      "      (conv): Conv2d(1024, 1024, kernel_size=(2, 2), stride=(2, 2), bias=False)\n",
      "    )\n",
      "  )\n",
      "  (up_block): Simple_UVitBlock(\n",
      "    (upsample): Upsample2D(\n",
      "      (norm): RMSNorm()\n",
      "      (conv): ConvTranspose2d(1024, 1024, kernel_size=(2, 2), stride=(2, 2), bias=False)\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from safetensors.torch import load_file\n",
    "import json\n",
    "from models.meissonic import Transformer2DModel\n",
    "\n",
    "model_path = \"/home/dongpeijie/.cache/huggingface/hub/models--MeissonFlow--Meissonic/snapshots/08ff13de62d55a6984806076d005089acc63f9ee/transformer/\"\n",
    "config_name = \"config.json\"\n",
    "file_name = \"diffusion_pytorch_model.safetensors\"\n",
    "\n",
    "loaded = load_file(model_path + file_name)\n",
    "\n",
    "# 打开一个包含JSON数据的文件\n",
    "with open(model_path + config_name, 'r') as file:\n",
    "    # 使用json.load()将文件中的JSON数据解码为Python对象\n",
    "    model_config_dict = json.load(file)\n",
    "\n",
    "meissonic_mdoel = Transformer2DModel(\n",
    "    patch_size=model_config_dict['patch_size'],\n",
    "    in_channels=model_config_dict['in_channels'],\n",
    "    num_layers=model_config_dict['num_layers'],\n",
    "    num_single_layers=model_config_dict['num_single_layers'],\n",
    "    attention_head_dim=model_config_dict['attention_head_dim'],\n",
    "    num_attention_heads=model_config_dict['num_attention_heads'],\n",
    "    joint_attention_dim=model_config_dict['joint_attention_dim'],\n",
    "    pooled_projection_dim= model_config_dict['pooled_projection_dim'],\n",
    "    guidance_embeds=model_config_dict['guidance_embeds'], # unused in our implementation\n",
    "    axes_dims_rope=tuple(model_config_dict['axes_dims_rope']),\n",
    "    vocab_size=model_config_dict['vocab_size'],\n",
    "    codebook_size=model_config_dict['codebook_size'],\n",
    "    downsample=model_config_dict['downsample'],\n",
    "    upsample=model_config_dict['upsample'],\n",
    ")\n",
    "\n",
    "meissonic_mdoel.load_state_dict(loaded)\n",
    "\n",
    "print(meissonic_mdoel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "# Set the manual seed for reproducibility\n",
    "torch.manual_seed(42)\n",
    "\n",
    "# Generate tensors based on the provided sizes and types\n",
    "hidden_states = torch.randint(0, 100, (2, 64, 64), dtype=torch.int64)\n",
    "micro_conds = torch.randn(2, 5, dtype=torch.float32)\n",
    "pooled_projections = torch.randn(2, 1024, dtype=torch.float32)\n",
    "encoder_hidden_states = torch.randn(2, 77, 1024, dtype=torch.float32)\n",
    "img_ids = torch.randint(0, 100, (1024, 3), dtype=torch.int64)\n",
    "txt_ids = torch.randn(77, 3, dtype=torch.float32)\n",
    "timestep = torch.randint(0, 10, (1,), dtype=torch.int64)\n",
    "\n",
    "# Print the generated tensors to verify\n",
    "# print(\"hidden_states:\", hidden_states)\n",
    "# print(\"micro_conds:\", micro_conds)\n",
    "# print(\"pooled_projections:\", pooled_projections)\n",
    "# print(\"encoder_hidden_states:\", encoder_hidden_states)\n",
    "# print(\"img_ids:\", img_ids)\n",
    "# print(\"txt_ids:\", txt_ids)\n",
    "# print(\"timestep:\", timestep)\n",
    "\n",
    "cur_res = meissonic_mdoel(\n",
    "    hidden_states = hidden_states,\n",
    "    micro_conds=micro_conds,\n",
    "    pooled_projections=pooled_projections,\n",
    "    encoder_hidden_states=encoder_hidden_states,\n",
    "    img_ids = img_ids,\n",
    "    txt_ids = txt_ids,\n",
    "    timestep = timestep,\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Tensor"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(cur_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 8192, 64, 64])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cur_res.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_842713/3036194550.py:2: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  true_res = torch.load(\"/home/dongpeijie/workspace/gg/Meissonic/test_res.pt\")\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "true_res = torch.load(\"/home/dongpeijie/workspace/gg/Meissonic/test_res.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 2.8369e-01, -1.2375e+00, -5.3806e+00,  ..., -3.4819e+00,\n",
       "           -4.0670e+00, -3.3877e+00],\n",
       "          [-2.9600e+00, -1.1494e+00, -9.7034e+00,  ..., -3.0984e+00,\n",
       "           -1.7382e+00, -3.2511e+00],\n",
       "          [-1.6342e+00,  6.3396e-01, -2.1614e+00,  ..., -4.5308e+00,\n",
       "            2.3254e+00, -2.8297e+00],\n",
       "          ...,\n",
       "          [-4.5445e+00, -2.3282e+00, -5.4029e+00,  ..., -6.6764e+00,\n",
       "            2.6510e-01, -1.9122e+00],\n",
       "          [-6.1303e-01, -2.4272e+00, -4.7205e+00,  ..., -2.7371e+00,\n",
       "           -6.0191e+00, -5.7930e+00],\n",
       "          [-6.3311e-01, -2.1056e+00, -1.8232e+00,  ...,  5.7443e-01,\n",
       "            2.8692e-01, -1.8012e+00]],\n",
       "\n",
       "         [[ 1.4578e+00, -1.2708e+00, -2.8397e+00,  ..., -7.5843e-01,\n",
       "           -1.7562e+00, -9.1217e-02],\n",
       "          [-2.4450e+00,  5.7523e-02, -4.1757e+00,  ..., -2.4705e+00,\n",
       "            1.8865e-01, -4.0500e-01],\n",
       "          [-3.3839e-01,  1.0405e+00, -7.0991e-01,  ...,  6.0070e-01,\n",
       "            3.1506e+00, -2.4566e-01],\n",
       "          ...,\n",
       "          [-2.5527e+00, -1.8289e+00, -2.9968e+00,  ..., -1.9658e+00,\n",
       "            8.1111e-02, -6.3732e-01],\n",
       "          [ 1.5277e+00, -5.1061e-01, -1.9020e+00,  ..., -1.6092e+00,\n",
       "           -3.7077e+00, -9.1595e-01],\n",
       "          [ 1.8600e-01, -6.8978e-01,  8.2091e-01,  ...,  6.3306e-01,\n",
       "            4.4220e+00, -1.9674e+00]],\n",
       "\n",
       "         [[-1.3447e+00, -2.5107e+00, -3.1962e+00,  ..., -3.9929e+00,\n",
       "           -2.8329e+00, -2.1886e+00],\n",
       "          [-1.8926e-01, -3.0744e-01, -2.4778e+00,  ..., -7.6008e+00,\n",
       "           -5.2844e-01, -2.9080e+00],\n",
       "          [-4.2397e-06,  1.0322e-01, -2.6643e+00,  ..., -8.4992e-02,\n",
       "           -4.3906e+00, -4.5985e+00],\n",
       "          ...,\n",
       "          [-3.8901e+00, -3.3211e+00, -1.9095e+00,  ..., -2.4703e+00,\n",
       "            9.9677e-01,  1.1264e+00],\n",
       "          [ 9.9851e-01, -6.1210e-01, -4.4066e+00,  ..., -1.7575e+00,\n",
       "           -4.8300e+00, -2.3406e+00],\n",
       "          [-1.4662e-01, -3.5549e-01, -1.6138e+00,  ..., -1.2592e+00,\n",
       "            2.2997e+00, -5.1713e+00]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 1.3931e+00,  6.3057e+00,  7.2207e+00,  ...,  1.2520e+00,\n",
       "            2.3157e+00,  7.3007e-01],\n",
       "          [ 2.6675e+00, -2.7546e-01,  2.6771e+00,  ...,  6.3396e+00,\n",
       "            1.0062e+00,  2.2980e+00],\n",
       "          [ 2.0984e+00,  2.9577e+00,  3.0466e+00,  ..., -5.6148e-01,\n",
       "            5.5100e-01,  6.1499e+00],\n",
       "          ...,\n",
       "          [ 5.2544e+00,  1.3380e+00,  3.5067e+00,  ...,  6.8607e-01,\n",
       "            2.4173e+00,  2.5298e+00],\n",
       "          [ 1.5433e+00,  2.2880e+00,  5.1108e+00,  ...,  3.5856e+00,\n",
       "            3.3221e+00,  4.4952e-01],\n",
       "          [ 1.6844e+00,  1.2417e+00,  2.5477e+00,  ...,  2.1621e+00,\n",
       "           -1.7824e+00,  4.9170e+00]],\n",
       "\n",
       "         [[ 3.6964e+00,  5.2578e+00,  7.9339e+00,  ...,  1.1906e+00,\n",
       "            4.5003e+00,  2.4179e+00],\n",
       "          [ 4.8952e+00,  9.2225e-01, -9.7831e-01,  ...,  4.6108e+00,\n",
       "            3.3152e+00,  4.1014e+00],\n",
       "          [ 3.8327e+00,  5.5375e+00,  4.3080e+00,  ...,  1.5091e+00,\n",
       "            3.2873e+00,  4.4905e+00],\n",
       "          ...,\n",
       "          [ 2.0629e+00,  1.2033e+00,  1.6773e+00,  ..., -1.9429e+00,\n",
       "            4.1567e+00,  2.4544e+00],\n",
       "          [ 3.1063e+00,  1.9850e+00,  5.1167e+00,  ...,  2.2039e+00,\n",
       "            4.9973e+00, -2.4433e-01],\n",
       "          [ 2.7603e+00,  1.4874e+00,  2.8107e+00,  ...,  1.7148e+00,\n",
       "            3.2013e-01,  5.3285e+00]],\n",
       "\n",
       "         [[ 1.0752e+00, -3.6909e+00, -1.0307e+00,  ..., -7.8230e-01,\n",
       "           -1.0676e+00,  8.9383e-01],\n",
       "          [-1.2368e+00,  8.4495e-01,  5.1956e-01,  ..., -2.9804e+00,\n",
       "           -3.6232e-01, -7.8727e-01],\n",
       "          [-8.7499e-01, -1.6317e+00, -9.9673e-01,  ...,  4.4757e+00,\n",
       "            2.5209e+00, -2.6243e+00],\n",
       "          ...,\n",
       "          [-1.8521e+00, -4.4657e-01, -1.5976e+00,  ...,  1.8721e+00,\n",
       "           -7.3255e-01, -3.2541e-01],\n",
       "          [ 1.5955e+00, -2.9060e-01, -2.2919e+00,  ..., -1.6728e+00,\n",
       "           -4.6535e-01,  2.1741e+00],\n",
       "          [ 4.3279e-01, -3.0146e-02, -4.2486e-01,  ..., -4.8102e-01,\n",
       "            5.4201e+00, -5.3299e-01]]],\n",
       "\n",
       "\n",
       "        [[[ 6.4904e-01, -2.7794e+00, -4.2601e+00,  ..., -3.4744e+00,\n",
       "           -7.1076e+00,  5.4553e+00],\n",
       "          [-1.9297e+00, -1.3619e-01, -9.8502e-02,  ..., -1.8758e+00,\n",
       "           -3.2768e+00, -1.1211e+00],\n",
       "          [-4.2955e+00,  5.2635e+00,  6.2904e-01,  ..., -4.9193e+00,\n",
       "           -6.0881e+00, -7.1501e+00],\n",
       "          ...,\n",
       "          [-4.5257e+00, -5.2411e+00, -4.3531e+00,  ...,  1.4749e+00,\n",
       "           -1.4293e+00, -3.4731e+00],\n",
       "          [-8.2702e-01, -2.2387e+00, -6.9634e+00,  ..., -2.7126e+00,\n",
       "           -3.2587e-01,  1.5172e+00],\n",
       "          [-2.1623e+00, -1.1458e+00, -5.1747e+00,  ..., -6.9246e-01,\n",
       "           -3.4658e+00, -2.7493e+00]],\n",
       "\n",
       "         [[-5.6817e-01, -2.4904e+00, -1.4954e+00,  ..., -2.6003e+00,\n",
       "           -3.1623e+00,  4.0189e+00],\n",
       "          [-2.6254e+00,  7.9034e-01,  2.1061e+00,  ..., -2.2880e-01,\n",
       "           -2.6187e+00,  2.1025e+00],\n",
       "          [-1.4175e+00,  6.2811e+00,  1.4567e+00,  ..., -1.1110e+00,\n",
       "           -4.0205e+00, -3.5389e+00],\n",
       "          ...,\n",
       "          [ 4.7369e-01, -1.0396e+00, -2.0605e+00,  ...,  2.6449e+00,\n",
       "            3.6172e-01, -6.5220e-01],\n",
       "          [-1.5018e+00, -1.0254e-02, -3.3405e+00,  ..., -1.2735e+00,\n",
       "            5.2641e-01,  4.4331e+00],\n",
       "          [-1.4096e+00, -1.3102e+00, -1.8172e+00,  ..., -4.3038e-01,\n",
       "           -1.3670e+00,  2.2092e-01]],\n",
       "\n",
       "         [[-3.6049e-01, -1.1850e+00, -1.7805e+00,  ..., -2.8893e+00,\n",
       "           -8.3347e+00, -3.0807e+00],\n",
       "          [-1.2037e+00, -2.1514e-01,  4.6734e-01,  ..., -7.4126e-01,\n",
       "           -8.8896e+00, -6.5471e-04],\n",
       "          [-4.9042e+00, -3.3430e+00, -3.8714e+00,  ..., -8.7866e+00,\n",
       "           -5.9451e+00, -4.9384e+00],\n",
       "          ...,\n",
       "          [-3.8044e+00, -3.5565e+00, -2.4975e+00,  ..., -4.7081e+00,\n",
       "           -1.6460e+00, -1.4118e+00],\n",
       "          [-5.9092e+00, -3.7441e+00, -6.6386e+00,  ...,  1.8747e+00,\n",
       "            1.5484e+00,  2.5945e+00],\n",
       "          [-3.4896e+00, -4.5170e+00, -7.2698e+00,  ...,  1.7681e-03,\n",
       "            3.0437e-01,  2.8116e+00]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 2.4559e+00,  5.6339e+00,  5.8752e+00,  ...,  5.8855e+00,\n",
       "            7.9009e+00,  3.2546e+00],\n",
       "          [ 3.1875e+00,  4.4554e-01,  1.0682e+00,  ...,  1.9817e+00,\n",
       "            1.2375e+01, -1.6082e-01],\n",
       "          [ 6.6627e+00,  2.8857e+00,  4.8532e+00,  ...,  5.4409e+00,\n",
       "           -6.0638e-01, -2.2801e+00],\n",
       "          ...,\n",
       "          [ 1.5099e+00,  8.6207e-01,  1.8558e+00,  ...,  3.1716e+00,\n",
       "            3.7179e+00,  9.0211e-01],\n",
       "          [ 4.5405e+00,  4.2958e+00,  7.3203e+00,  ...,  1.7869e-01,\n",
       "            1.5659e+00,  3.8548e-01],\n",
       "          [ 2.0591e+00,  4.6672e+00,  7.8843e+00,  ...,  2.3295e+00,\n",
       "            1.0892e+00,  2.5780e-02]],\n",
       "\n",
       "         [[ 3.5092e+00,  4.6323e+00,  4.1877e+00,  ...,  7.2885e+00,\n",
       "            3.0406e+00,  1.3107e+01],\n",
       "          [ 2.0582e+00,  3.5956e-01,  4.5604e+00,  ...,  4.1620e+00,\n",
       "            9.2146e+00,  2.8413e+00],\n",
       "          [ 4.2612e+00,  7.9684e+00,  6.0910e+00,  ...,  2.2650e+00,\n",
       "           -1.4633e+00, -2.4926e+00],\n",
       "          ...,\n",
       "          [ 2.3655e+00, -1.1146e-01,  1.5715e+00,  ...,  6.0058e+00,\n",
       "            5.4539e+00,  7.1749e-01],\n",
       "          [ 4.3858e+00,  4.7016e+00,  3.1981e+00,  ...,  4.8542e-01,\n",
       "            5.3981e+00,  3.2254e+00],\n",
       "          [ 1.4610e+00,  3.4045e+00,  2.5838e+00,  ...,  3.2965e+00,\n",
       "            2.0460e+00,  1.4771e+00]],\n",
       "\n",
       "         [[ 1.7597e-01, -2.3389e+00, -1.2321e+00,  ..., -3.7142e+00,\n",
       "           -1.7271e+00, -9.2245e-01],\n",
       "          [-2.0346e+00,  2.2258e+00,  3.2026e+00,  ...,  1.3210e+00,\n",
       "           -3.4711e+00,  1.8908e+00],\n",
       "          [-2.5589e+00,  1.9744e+00, -5.9139e-01,  ..., -1.9156e+00,\n",
       "            2.3838e+00,  6.5960e+00],\n",
       "          ...,\n",
       "          [ 1.4393e+00,  1.2902e+00, -1.0672e+00,  ...,  1.0657e+00,\n",
       "           -9.9649e-01, -1.3910e-01],\n",
       "          [-2.0092e+00, -1.0020e-01, -2.8013e+00,  ...,  7.2120e-01,\n",
       "            1.1099e-01,  1.6468e+00],\n",
       "          [ 1.0724e-01, -2.2973e+00, -2.4348e+00,  ..., -7.7205e-01,\n",
       "            4.1957e-02,  1.6712e+00]]]], grad_fn=<ConvolutionBackward0>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cur_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.equal(true_res, cur_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 8192, 64, 64])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "true_res.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1]"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l  = [1, 2, 3]\n",
    "l[0: 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fluxgym",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}