riskyhomo commited on
Commit
6817d94
1 Parent(s): a280ee2

Upload 2 files

Browse files
Files changed (2) hide show
  1. decoder.ipynb +133 -0
  2. encoder.ipynb +92 -0
decoder.ipynb ADDED
@@ -0,0 +1,133 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "nbformat": 4,
3
+ "nbformat_minor": 0,
4
+ "metadata": {
5
+ "colab": {
6
+ "provenance": []
7
+ },
8
+ "kernelspec": {
9
+ "name": "python3",
10
+ "display_name": "Python 3"
11
+ },
12
+ "language_info": {
13
+ "name": "python"
14
+ }
15
+ },
16
+ "cells": [
17
+ {
18
+ "cell_type": "code",
19
+ "execution_count": 8,
20
+ "metadata": {
21
+ "id": "ClbDA89uqYc2",
22
+ "colab": {
23
+ "base_uri": "https://localhost:8080/",
24
+ "height": 383
25
+ },
26
+ "outputId": "0dac9131-1d0c-416b-98fb-9dcc7e0dcf5b"
27
+ },
28
+ "outputs": [
29
+ {
30
+ "output_type": "error",
31
+ "ename": "ModuleNotFoundError",
32
+ "evalue": "No module named 'attention'",
33
+ "traceback": [
34
+ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
35
+ "\u001b[0;31mModuleNotFoundError\u001b[0m Traceback (most recent call last)",
36
+ "\u001b[0;32m<ipython-input-8-58da732233b1>\u001b[0m in \u001b[0;36m<cell line: 4>\u001b[0;34m()\u001b[0m\n\u001b[1;32m 2\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mfunctional\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mattention\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mSelfAttention\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 6\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mVAE_AttentionBlock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mModule\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
37
+ "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'attention'",
38
+ "",
39
+ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0;32m\nNOTE: If your import is failing due to a missing package, you can\nmanually install dependencies using either !pip or !apt.\n\nTo view examples of installing some common dependencies, click the\n\"Open Examples\" button below.\n\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n"
40
+ ],
41
+ "errorDetails": {
42
+ "actions": [
43
+ {
44
+ "action": "open_url",
45
+ "actionText": "Open Examples",
46
+ "url": "/notebooks/snippets/importing_libraries.ipynb"
47
+ }
48
+ ]
49
+ }
50
+ }
51
+ ],
52
+ "source": [
53
+ "import torch\n",
54
+ "from torch import nn\n",
55
+ "from torch.nn import functional as F\n",
56
+ "from attention import SelfAttention\n",
57
+ "\n",
58
+ "class VAE_AttentionBlock(nn.Module):\n",
59
+ " def __init__(self, channels):\n",
60
+ " super.__init__()\n",
61
+ " self.groupnorm = nn.GroupNorm(32, channels)\n",
62
+ " self.attention = Attention(1, channels)\n",
63
+ " def forward(self, x):\n",
64
+ " residue=x\n",
65
+ " x=self.groupnorm(x)\n",
66
+ " n,c,h,w =x.shape\n",
67
+ " x=x.view(n,c,h*w)\n",
68
+ " x=x.transpose(-1,-2)\n",
69
+ " x=self.attention(x)\n",
70
+ " x=x.view((n,c,h,w))\n",
71
+ " x+=residue\n",
72
+ " return x\n",
73
+ "class VAE_ResidualBlock(nn.Module):\n",
74
+ " def __init__(self, in_channels, out_channels):\n",
75
+ " super.__init__()\n",
76
+ " self.groupnorm_1= nn.GroupNorm(32, in_channels)\n",
77
+ " self.conv_1=nn.Conv2d(in_channels, out_channels , kernel_size=3, padding=1)\n",
78
+ "\n",
79
+ " self.groupnorm_2= nn.GroupNorm(32, in_channels)\n",
80
+ " self.conv_2= nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)\n",
81
+ "\n",
82
+ " if in_channels==out_channels:\n",
83
+ " self.residual_layer=nn.Identity()\n",
84
+ " else:\n",
85
+ " self.conv_2= nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0)\n",
86
+ " def forward(self, x):\n",
87
+ " residue=x\n",
88
+ " x=self.groupnorm_1(x)\n",
89
+ " x=F.silu(x)\n",
90
+ " x=self.conv_2(x)\n",
91
+ " return x+self.residual_layer(residue)\n",
92
+ "class VAE_Decoder(nn.Sequential):\n",
93
+ " def __init__(self):\n",
94
+ " super.__init__(\n",
95
+ " nn.Conv2d(4,4, kernel_size=1, padding=0),\n",
96
+ " nn.Conv2d(4, 512, kernel_size=3, padding=1),\n",
97
+ " VAE_ResidualBlock(512, 512),\n",
98
+ " VAE_AttentionBlock(512),\n",
99
+ " VAE_ResidualBlock(512, 512),\n",
100
+ " VAE_ResidualBlock(512, 512),\n",
101
+ " VAE_ResidualBlock(512, 512),\n",
102
+ " VAE_ResidualBlock(512, 512),\n",
103
+ " VAE_ResidualBlock(512, 512),\n",
104
+ "\n",
105
+ " nn.Upsample(scale_factor=2),\n",
106
+ " nn.Conv2d(512, 512,kernel_size=3, padding=1),\n",
107
+ " VAE_ResidualBlock(512, 512),\n",
108
+ " VAE_ResidualBlock(512, 512),\n",
109
+ " VAE_ResidualBlock(512, 512),\n",
110
+ " nn.Upsample(scale_factor=2),\n",
111
+ " VAE_ResidualBlock(512, 256),\n",
112
+ " VAE_ResidualBlock(256, 256),\n",
113
+ " VAE_ResidualBlock(256, 256),\n",
114
+ " nn.Upsample(scale_factor=2),\n",
115
+ " nn.Conv2d(256, 256, kernel_size=3, padding=1),\n",
116
+ " VAE_ResidualBlock(256, 128),\n",
117
+ " VAE_ResidualBlock(128, 128),\n",
118
+ " VAE_ResidualBlock(128, 128),\n",
119
+ " nn.GroupNorm(32, 128),\n",
120
+ " nn.SiLU(),\n",
121
+ " nn.Conv2d(128, 3, kernel_size=3,padding=1)\n",
122
+ " )\n",
123
+ " def forward(self, x):\n",
124
+ " x/=0.18125\n",
125
+ " for module in self:\n",
126
+ " x=module(x)\n",
127
+ " return x\n",
128
+ "\n",
129
+ "\n"
130
+ ]
131
+ }
132
+ ]
133
+ }
encoder.ipynb ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "nbformat": 4,
3
+ "nbformat_minor": 0,
4
+ "metadata": {
5
+ "colab": {
6
+ "provenance": []
7
+ },
8
+ "kernelspec": {
9
+ "name": "python3",
10
+ "display_name": "Python 3"
11
+ },
12
+ "language_info": {
13
+ "name": "python"
14
+ }
15
+ },
16
+ "cells": [
17
+ {
18
+ "cell_type": "code",
19
+ "execution_count": 1,
20
+ "metadata": {
21
+ "id": "OaBIb0WNma-U",
22
+ "colab": {
23
+ "base_uri": "https://localhost:8080/",
24
+ "height": 383
25
+ },
26
+ "outputId": "33f26e87-f2b5-4de8-dfd8-dca7e8973a01"
27
+ },
28
+ "outputs": [
29
+ {
30
+ "output_type": "error",
31
+ "ename": "ModuleNotFoundError",
32
+ "evalue": "No module named 'decoder'",
33
+ "traceback": [
34
+ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
35
+ "\u001b[0;31mModuleNotFoundError\u001b[0m Traceback (most recent call last)",
36
+ "\u001b[0;32m<ipython-input-1-140de5a90328>\u001b[0m in \u001b[0;36m<cell line: 4>\u001b[0;34m()\u001b[0m\n\u001b[1;32m 2\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mfunctional\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mdecoder\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mVAE_AttentionBlock\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mVAE_ResidualBlock\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 6\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mVAE_Encoder\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSequential\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
37
+ "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'decoder'",
38
+ "",
39
+ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0;32m\nNOTE: If your import is failing due to a missing package, you can\nmanually install dependencies using either !pip or !apt.\n\nTo view examples of installing some common dependencies, click the\n\"Open Examples\" button below.\n\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n"
40
+ ],
41
+ "errorDetails": {
42
+ "actions": [
43
+ {
44
+ "action": "open_url",
45
+ "actionText": "Open Examples",
46
+ "url": "/notebooks/snippets/importing_libraries.ipynb"
47
+ }
48
+ ]
49
+ }
50
+ }
51
+ ],
52
+ "source": [
53
+ "import torch\n",
54
+ "from torch import nn\n",
55
+ "from torch.nn import functional as F\n",
56
+ "from decoder import VAE_AttentionBlock, VAE_ResidualBlock\n",
57
+ "\n",
58
+ "class VAE_Encoder(nn.Sequential):\n",
59
+ " def __init__(self):\n",
60
+ " super().__init__(\n",
61
+ " nn.Conv2d(3, 128, kernel_size=3, padding=1),\n",
62
+ " VAE_ResidualBlock(128, 128),\n",
63
+ " VAE_ResidualBlock(128, 128),\n",
64
+ " nn.Conv2d(128, 128, kernel_size=3,stride=2, padding=1),\n",
65
+ " VAE_ResidualBlock(128,256),\n",
66
+ " VAE_ResidualBlock(256, 256),\n",
67
+ " nn.Conv2d(3, 128, kernel_size=3,stride=2, padding=1),\n",
68
+ " VAE_ResidualBlock(256, 512),\n",
69
+ " VAE_ResidualBlock(512, 512),\n",
70
+ " VAE_AttentionBlock(512),\n",
71
+ " VAE_ResidualBlock(512, 512),\n",
72
+ " nn.GroupNorm(32, 512),\n",
73
+ " nn.SiLU(),\n",
74
+ " nn.Conv2d(512, 8, kernel_size=3, padding=1),\n",
75
+ " nn.Conv2d(8,8,kernel_size=3, padding=0),\n",
76
+ " )\n",
77
+ " def forward(self, x , noise):\n",
78
+ " for module in self:\n",
79
+ " if getattr(module, \"stride\", None) ==(2,2):\n",
80
+ " x= F.pad(x, (0,1,0,1))\n",
81
+ " x= module(x)\n",
82
+ " mean , log_variance = torch.chunk(x, 2, dim=1)\n",
83
+ " log_variance= torch.clamp(log_variance, -20, 30)\n",
84
+ " variance= log_variance.exp()\n",
85
+ " stdev = variance.sqrt()\n",
86
+ " x= stdev*noise +mean\n",
87
+ " x*=0.18215\n",
88
+ " return x\n"
89
+ ]
90
+ }
91
+ ]
92
+ }