├── .gitignore
├── .idea
    ├── inspectionProfiles
    │   └── Project_Default.xml
    ├── misc.xml
    ├── modules.xml
    ├── tensorflow-mnist-VAE.iml
    ├── vcs.xml
    └── workspace.xml
├── README.md
├── __pycache__
    ├── mnist_data.cpython-36.pyc
    ├── plot_utils.cpython-36.pyc
    └── vae.cpython-36.pyc
├── mnist_data.py
├── plot_utils.py
├── res
    ├── denoising.jpg
    ├── dim_10_sample.jpg
    ├── dim_20_sample.jpg
    ├── dim_2_cluster.jpg
    ├── dim_2_mf.jpg
    ├── dim_2_sample.jpg
    ├── dim_5_sample.jpg
    ├── input.jpg
    └── input_noise.jpg
├── run_main.py
├── test.py
├── vae.py
└── vae2.py


/.gitignore:
--------------------------------------------------------------------------------
1 | results
2 | data
3 | __pycache__
4 | 


--------------------------------------------------------------------------------
/.idea/inspectionProfiles/Project_Default.xml:
--------------------------------------------------------------------------------
1 | <component name="InspectionProjectProfileManager">
2 |   <profile version="1.0">
3 |     <option name="myName" value="Project Default" />
4 |     <inspection_tool class="PyPep8Inspection" enabled="false" level="WEAK WARNING" enabled_by_default="false" />
5 |   </profile>
6 | </component>


--------------------------------------------------------------------------------
/.idea/misc.xml:
--------------------------------------------------------------------------------
1 | <?xml version="1.0" encoding="UTF-8"?>
2 | <project version="4">
3 |   <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.6 (FAUL)" project-jdk-type="Python SDK" />
4 | </project>


--------------------------------------------------------------------------------
/.idea/modules.xml:
--------------------------------------------------------------------------------
1 | <?xml version="1.0" encoding="UTF-8"?>
2 | <project version="4">
3 |   <component name="ProjectModuleManager">
4 |     <modules>
5 |       <module fileurl="file://$PROJECT_DIR$/.idea/tensorflow-mnist-VAE.iml" filepath="$PROJECT_DIR$/.idea/tensorflow-mnist-VAE.iml" />
6 |     </modules>
7 |   </component>
8 | </project>


--------------------------------------------------------------------------------
/.idea/tensorflow-mnist-VAE.iml:
--------------------------------------------------------------------------------
 1 | <?xml version="1.0" encoding="UTF-8"?>
 2 | <module type="PYTHON_MODULE" version="4">
 3 |   <component name="NewModuleRootManager">
 4 |     <content url="file://$MODULE_DIR$" />
 5 |     <orderEntry type="inheritedJdk" />
 6 |     <orderEntry type="sourceFolder" forTests="false" />
 7 |   </component>
 8 |   <component name="TestRunnerService">
 9 |     <option name="projectConfiguration" value="pytest" />
10 |     <option name="PROJECT_TEST_RUNNER" value="pytest" />
11 |   </component>
12 | </module>


--------------------------------------------------------------------------------
/.idea/vcs.xml:
--------------------------------------------------------------------------------
1 | <?xml version="1.0" encoding="UTF-8"?>
2 | <project version="4">
3 |   <component name="VcsDirectoryMappings">
4 |     <mapping directory="$PROJECT_DIR$" vcs="Git" />
5 |   </component>
6 | </project>


--------------------------------------------------------------------------------
/.idea/workspace.xml:
--------------------------------------------------------------------------------
  1 | <?xml version="1.0" encoding="UTF-8"?>
  2 | <project version="4">
  3 |   <component name="ChangeListManager">
  4 |     <list default="true" id="b7c2c026-5f3b-469a-ae78-61f458e8da2a" name="Default Changelist" comment="">
  5 |       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
  6 |       <change beforePath="$PROJECT_DIR$/run_main.py" beforeDir="false" afterPath="$PROJECT_DIR$/run_main.py" afterDir="false" />
  7 |       <change beforePath="$PROJECT_DIR$/vae.py" beforeDir="false" afterPath="$PROJECT_DIR$/vae.py" afterDir="false" />
  8 |     </list>
  9 |     <option name="EXCLUDED_CONVERTED_TO_IGNORED" value="true" />
 10 |     <option name="SHOW_DIALOG" value="false" />
 11 |     <option name="HIGHLIGHT_CONFLICTS" value="true" />
 12 |     <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
 13 |     <option name="LAST_RESOLUTION" value="IGNORE" />
 14 |   </component>
 15 |   <component name="FUSProjectUsageTrigger">
 16 |     <session id="964659601">
 17 |       <usages-collector id="statistics.lifecycle.project">
 18 |         <counts>
 19 |           <entry key="project.closed" value="1" />
 20 |           <entry key="project.open.time.0" value="1" />
 21 |           <entry key="project.opened" value="1" />
 22 |         </counts>
 23 |       </usages-collector>
 24 |       <usages-collector id="statistics.file.extensions.open">
 25 |         <counts>
 26 |           <entry key="py" value="4" />
 27 |         </counts>
 28 |       </usages-collector>
 29 |       <usages-collector id="statistics.file.types.open">
 30 |         <counts>
 31 |           <entry key="Python" value="4" />
 32 |         </counts>
 33 |       </usages-collector>
 34 |     </session>
 35 |     <session id="-2093924913">
 36 |       <usages-collector id="statistics.lifecycle.project">
 37 |         <counts>
 38 |           <entry key="project.closed" value="2" />
 39 |           <entry key="project.open.time.0" value="1" />
 40 |           <entry key="project.open.time.7" value="1" />
 41 |           <entry key="project.open.time.8" value="1" />
 42 |           <entry key="project.opened" value="3" />
 43 |         </counts>
 44 |       </usages-collector>
 45 |       <usages-collector id="statistics.file.extensions.edit">
 46 |         <counts>
 47 |           <entry key="py" value="2575" />
 48 |         </counts>
 49 |       </usages-collector>
 50 |       <usages-collector id="statistics.file.types.edit">
 51 |         <counts>
 52 |           <entry key="Python" value="2575" />
 53 |         </counts>
 54 |       </usages-collector>
 55 |       <usages-collector id="statistics.file.extensions.open">
 56 |         <counts>
 57 |           <entry key="py" value="6" />
 58 |         </counts>
 59 |       </usages-collector>
 60 |       <usages-collector id="statistics.file.types.open">
 61 |         <counts>
 62 |           <entry key="Python" value="6" />
 63 |         </counts>
 64 |       </usages-collector>
 65 |     </session>
 66 |   </component>
 67 |   <component name="FileEditorManager">
 68 |     <leaf SIDE_TABS_SIZE_LIMIT_KEY="300">
 69 |       <file pinned="false" current-in-tab="true">
 70 |         <entry file="file://$PROJECT_DIR$/run_main.py">
 71 |           <provider selected="true" editor-type-id="text-editor">
 72 |             <state relative-caret-position="301">
 73 |               <caret line="220" column="20" lean-forward="true" selection-start-line="220" selection-start-column="20" selection-end-line="220" selection-end-column="20" />
 74 |               <folding>
 75 |                 <element signature="e#0#12#0" expanded="true" />
 76 |                 <element signature="e#2540#5351#0" />
 77 |               </folding>
 78 |             </state>
 79 |           </provider>
 80 |         </entry>
 81 |       </file>
 82 |       <file pinned="false" current-in-tab="false">
 83 |         <entry file="file://$PROJECT_DIR$/plot_utils.py">
 84 |           <provider selected="true" editor-type-id="text-editor">
 85 |             <state relative-caret-position="2420">
 86 |               <caret line="113" column="8" selection-start-line="113" selection-start-column="8" selection-end-line="113" selection-end-column="8" />
 87 |               <folding>
 88 |                 <element signature="e#0#18#0" expanded="true" />
 89 |               </folding>
 90 |             </state>
 91 |           </provider>
 92 |         </entry>
 93 |       </file>
 94 |       <file pinned="false" current-in-tab="false">
 95 |         <entry file="file://$PROJECT_DIR$/mnist_data.py">
 96 |           <provider selected="true" editor-type-id="text-editor">
 97 |             <state relative-caret-position="2310">
 98 |               <caret line="117" column="4" selection-start-line="117" selection-start-column="4" selection-end-line="117" selection-end-column="4" />
 99 |             </state>
100 |           </provider>
101 |         </entry>
102 |       </file>
103 |       <file pinned="false" current-in-tab="false">
104 |         <entry file="file://$PROJECT_DIR$/vae.py">
105 |           <provider selected="true" editor-type-id="text-editor">
106 |             <state relative-caret-position="726">
107 |               <caret line="66" column="21" selection-start-line="66" selection-start-column="21" selection-end-line="66" selection-end-column="21" />
108 |             </state>
109 |           </provider>
110 |         </entry>
111 |       </file>
112 |       <file pinned="false" current-in-tab="false">
113 |         <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/torch/nn/modules/dropout.py">
114 |           <provider selected="true" editor-type-id="text-editor">
115 |             <state relative-caret-position="396">
116 |               <caret line="19" column="6" selection-start-line="19" selection-start-column="6" selection-end-line="19" selection-end-column="6" />
117 |             </state>
118 |           </provider>
119 |         </entry>
120 |       </file>
121 |       <file pinned="false" current-in-tab="false">
122 |         <entry file="file://$PROJECT_DIR$/test.py">
123 |           <provider selected="true" editor-type-id="text-editor">
124 |             <state relative-caret-position="550">
125 |               <caret line="27" column="30" selection-start-line="27" selection-start-column="30" selection-end-line="27" selection-end-column="30" />
126 |               <folding>
127 |                 <element signature="e#0#13#0" expanded="true" />
128 |                 <marker date="1542183498758" expanded="true" signature="128:648" ph="..." />
129 |               </folding>
130 |             </state>
131 |           </provider>
132 |         </entry>
133 |       </file>
134 |       <file pinned="false" current-in-tab="false">
135 |         <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/torch/nn/functional.py">
136 |           <provider selected="true" editor-type-id="text-editor">
137 |             <state relative-caret-position="18040">
138 |               <caret line="833" selection-start-line="833" selection-end-line="833" />
139 |             </state>
140 |           </provider>
141 |         </entry>
142 |       </file>
143 |       <file pinned="false" current-in-tab="false">
144 |         <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/torch/nn/init.py">
145 |           <provider selected="true" editor-type-id="text-editor">
146 |             <state relative-caret-position="6402">
147 |               <caret line="295" column="4" selection-start-line="295" selection-start-column="4" selection-end-line="295" selection-end-column="4" />
148 |             </state>
149 |           </provider>
150 |         </entry>
151 |       </file>
152 |       <file pinned="false" current-in-tab="false">
153 |         <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/torch/nn/modules/linear.py">
154 |           <provider selected="true" editor-type-id="text-editor">
155 |             <state relative-caret-position="682">
156 |               <caret line="36" column="8" selection-start-line="36" selection-start-column="8" selection-end-line="36" selection-end-column="8" />
157 |             </state>
158 |           </provider>
159 |         </entry>
160 |       </file>
161 |       <file pinned="false" current-in-tab="false">
162 |         <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/tensorflow/python/ops/gen_nn_ops.py">
163 |           <provider selected="true" editor-type-id="text-editor">
164 |             <state relative-caret-position="159236">
165 |               <caret line="7256" column="4" selection-start-line="7256" selection-start-column="4" selection-end-line="7256" selection-end-column="4" />
166 |             </state>
167 |           </provider>
168 |         </entry>
169 |       </file>
170 |     </leaf>
171 |   </component>
172 |   <component name="FileTemplateManagerImpl">
173 |     <option name="RECENT_TEMPLATES">
174 |       <list>
175 |         <option value="Python Script" />
176 |       </list>
177 |     </option>
178 |   </component>
179 |   <component name="FindInProjectRecents">
180 |     <findStrings>
181 |       <find>mnist_data</find>
182 |       <find>train_total_data</find>
183 |     </findStrings>
184 |   </component>
185 |   <component name="Git.Settings">
186 |     <option name="RECENT_GIT_ROOT_PATH" value="$PROJECT_DIR$" />
187 |   </component>
188 |   <component name="IdeDocumentHistory">
189 |     <option name="CHANGED_PATHS">
190 |       <list>
191 |         <option value="$PROJECT_DIR$/test.py" />
192 |         <option value="$PROJECT_DIR$/vae.py" />
193 |         <option value="$PROJECT_DIR$/run_main.py" />
194 |       </list>
195 |     </option>
196 |   </component>
197 |   <component name="ProjectFrameBounds" extendedState="6">
198 |     <option name="x" value="2027" />
199 |     <option name="y" value="163" />
200 |     <option name="width" value="1853" />
201 |     <option name="height" value="1076" />
202 |   </component>
203 |   <component name="ProjectInspectionProfilesVisibleTreeState">
204 |     <entry key="Project Default">
205 |       <profile-state>
206 |         <expanded-state>
207 |           <State />
208 |           <State>
209 |             <id>Python</id>
210 |           </State>
211 |         </expanded-state>
212 |         <selected-state>
213 |           <State>
214 |             <id>PyPep8Inspection</id>
215 |           </State>
216 |         </selected-state>
217 |       </profile-state>
218 |     </entry>
219 |   </component>
220 |   <component name="ProjectLevelVcsManager">
221 |     <ConfirmationsSetting value="1" id="Add" />
222 |   </component>
223 |   <component name="ProjectView">
224 |     <navigator proportions="" version="1">
225 |       <foldersAlwaysOnTop value="true" />
226 |     </navigator>
227 |     <panes>
228 |       <pane id="Scope" />
229 |       <pane id="ProjectPane">
230 |         <subPane>
231 |           <expand>
232 |             <path>
233 |               <item name="pytorch-mnist-VAE" type="b2602c69:ProjectViewProjectNode" />
234 |               <item name="pytorch-mnist-VAE" type="462c0819:PsiDirectoryNode" />
235 |             </path>
236 |           </expand>
237 |           <select />
238 |         </subPane>
239 |       </pane>
240 |     </panes>
241 |   </component>
242 |   <component name="PropertiesComponent">
243 |     <property name="settings.editor.selected.configurable" value="Errors" />
244 |   </component>
245 |   <component name="RunDashboard">
246 |     <option name="ruleStates">
247 |       <list>
248 |         <RuleState>
249 |           <option name="name" value="ConfigurationTypeDashboardGroupingRule" />
250 |         </RuleState>
251 |         <RuleState>
252 |           <option name="name" value="StatusDashboardGroupingRule" />
253 |         </RuleState>
254 |       </list>
255 |     </option>
256 |   </component>
257 |   <component name="RunManager">
258 |     <configuration name="test" type="PythonConfigurationType" factoryName="Python" temporary="true">
259 |       <module name="tensorflow-mnist-VAE" />
260 |       <option name="INTERPRETER_OPTIONS" value="" />
261 |       <option name="PARENT_ENVS" value="true" />
262 |       <envs>
263 |         <env name="PYTHONUNBUFFERED" value="1" />
264 |       </envs>
265 |       <option name="SDK_HOME" value="" />
266 |       <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
267 |       <option name="IS_MODULE_SDK" value="true" />
268 |       <option name="ADD_CONTENT_ROOTS" value="true" />
269 |       <option name="ADD_SOURCE_ROOTS" value="true" />
270 |       <option name="SCRIPT_NAME" value="$PROJECT_DIR$/test.py" />
271 |       <option name="PARAMETERS" value="" />
272 |       <option name="SHOW_COMMAND_LINE" value="false" />
273 |       <option name="EMULATE_TERMINAL" value="false" />
274 |       <option name="MODULE_MODE" value="false" />
275 |       <option name="REDIRECT_INPUT" value="false" />
276 |       <option name="INPUT_FILE" value="" />
277 |       <method v="2" />
278 |     </configuration>
279 |     <recent_temporary>
280 |       <list>
281 |         <item itemvalue="Python.test" />
282 |       </list>
283 |     </recent_temporary>
284 |   </component>
285 |   <component name="SvnConfiguration">
286 |     <configuration />
287 |   </component>
288 |   <component name="TaskManager">
289 |     <task active="true" id="Default" summary="Default task">
290 |       <changelist id="b7c2c026-5f3b-469a-ae78-61f458e8da2a" name="Default Changelist" comment="" />
291 |       <created>1542171285891</created>
292 |       <option name="number" value="Default" />
293 |       <option name="presentableId" value="Default" />
294 |       <updated>1542171285891</updated>
295 |     </task>
296 |     <servers />
297 |   </component>
298 |   <component name="ToolWindowManager">
299 |     <frame x="1984" y="-11" width="1857" height="1092" extended-state="6" />
300 |     <editor active="true" />
301 |     <layout>
302 |       <window_info active="true" content_ui="combo" id="Project" order="0" visible="true" weight="0.13864629" />
303 |       <window_info id="Structure" order="1" side_tool="true" weight="0.25" />
304 |       <window_info id="Favorites" order="2" side_tool="true" />
305 |       <window_info anchor="bottom" id="Message" order="0" />
306 |       <window_info anchor="bottom" id="Find" order="1" />
307 |       <window_info anchor="bottom" id="Run" order="2" weight="0.3294979" />
308 |       <window_info anchor="bottom" id="Debug" order="3" weight="0.4" />
309 |       <window_info anchor="bottom" id="Cvs" order="4" weight="0.25" />
310 |       <window_info anchor="bottom" id="Inspection" order="5" weight="0.4" />
311 |       <window_info anchor="bottom" id="TODO" order="6" />
312 |       <window_info anchor="bottom" id="Version Control" order="7" show_stripe_button="false" />
313 |       <window_info anchor="bottom" id="Terminal" order="8" />
314 |       <window_info anchor="bottom" id="Event Log" order="9" side_tool="true" />
315 |       <window_info anchor="bottom" id="Python Console" order="10" />
316 |       <window_info anchor="right" id="Commander" internal_type="SLIDING" order="0" type="SLIDING" weight="0.4" />
317 |       <window_info anchor="right" id="Ant Build" order="1" weight="0.25" />
318 |       <window_info anchor="right" content_ui="combo" id="Hierarchy" order="2" weight="0.25" />
319 |     </layout>
320 |   </component>
321 |   <component name="VcsContentAnnotationSettings">
322 |     <option name="myLimit" value="2678400000" />
323 |   </component>
324 |   <component name="XDebuggerManager">
325 |     <breakpoint-manager>
326 |       <breakpoints>
327 |         <line-breakpoint enabled="true" suspend="THREAD" type="python-line">
328 |           <url>file://$PROJECT_DIR$/run_main.py</url>
329 |           <line>214</line>
330 |           <option name="timeStamp" value="1" />
331 |         </line-breakpoint>
332 |       </breakpoints>
333 |     </breakpoint-manager>
334 |   </component>
335 |   <component name="editorHistoryManager">
336 |     <entry file="file://$PROJECT_DIR$/vae2.py">
337 |       <provider selected="true" editor-type-id="text-editor">
338 |         <state relative-caret-position="706">
339 |           <caret line="86" column="14" selection-start-line="86" selection-start-column="14" selection-end-line="86" selection-end-column="14" />
340 |         </state>
341 |       </provider>
342 |     </entry>
343 |     <entry file="file://$PROJECT_DIR$/plot_utils.py">
344 |       <provider selected="true" editor-type-id="text-editor">
345 |         <state relative-caret-position="2420">
346 |           <caret line="113" column="8" selection-start-line="113" selection-start-column="8" selection-end-line="113" selection-end-column="8" />
347 |           <folding>
348 |             <element signature="e#0#18#0" expanded="true" />
349 |           </folding>
350 |         </state>
351 |       </provider>
352 |     </entry>
353 |     <entry file="file://$PROJECT_DIR$/mnist_data.py">
354 |       <provider selected="true" editor-type-id="text-editor">
355 |         <state relative-caret-position="2310">
356 |           <caret line="117" column="4" selection-start-line="117" selection-start-column="4" selection-end-line="117" selection-end-column="4" />
357 |         </state>
358 |       </provider>
359 |     </entry>
360 |     <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/torch/nn/modules/dropout.py">
361 |       <provider selected="true" editor-type-id="text-editor">
362 |         <state relative-caret-position="396">
363 |           <caret line="19" column="6" selection-start-line="19" selection-start-column="6" selection-end-line="19" selection-end-column="6" />
364 |         </state>
365 |       </provider>
366 |     </entry>
367 |     <entry file="file://$PROJECT_DIR$/test.py">
368 |       <provider selected="true" editor-type-id="text-editor">
369 |         <state relative-caret-position="550">
370 |           <caret line="27" column="30" selection-start-line="27" selection-start-column="30" selection-end-line="27" selection-end-column="30" />
371 |           <folding>
372 |             <element signature="e#0#13#0" expanded="true" />
373 |             <marker date="1542183498758" expanded="true" signature="128:648" ph="..." />
374 |           </folding>
375 |         </state>
376 |       </provider>
377 |     </entry>
378 |     <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/torch/nn/functional.py">
379 |       <provider selected="true" editor-type-id="text-editor">
380 |         <state relative-caret-position="18040">
381 |           <caret line="833" selection-start-line="833" selection-end-line="833" />
382 |         </state>
383 |       </provider>
384 |     </entry>
385 |     <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/torch/nn/init.py">
386 |       <provider selected="true" editor-type-id="text-editor">
387 |         <state relative-caret-position="6402">
388 |           <caret line="295" column="4" selection-start-line="295" selection-start-column="4" selection-end-line="295" selection-end-column="4" />
389 |         </state>
390 |       </provider>
391 |     </entry>
392 |     <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/torch/nn/modules/linear.py">
393 |       <provider selected="true" editor-type-id="text-editor">
394 |         <state relative-caret-position="682">
395 |           <caret line="36" column="8" selection-start-line="36" selection-start-column="8" selection-end-line="36" selection-end-column="8" />
396 |         </state>
397 |       </provider>
398 |     </entry>
399 |     <entry file="file://$USER_HOME$/conda/lib/python3.6/site-packages/tensorflow/python/ops/gen_nn_ops.py">
400 |       <provider selected="true" editor-type-id="text-editor">
401 |         <state relative-caret-position="159236">
402 |           <caret line="7256" column="4" selection-start-line="7256" selection-start-column="4" selection-end-line="7256" selection-end-column="4" />
403 |         </state>
404 |       </provider>
405 |     </entry>
406 |     <entry file="file://$PROJECT_DIR$/vae.py">
407 |       <provider selected="true" editor-type-id="text-editor">
408 |         <state relative-caret-position="726">
409 |           <caret line="66" column="21" selection-start-line="66" selection-start-column="21" selection-end-line="66" selection-end-column="21" />
410 |         </state>
411 |       </provider>
412 |     </entry>
413 |     <entry file="file://$PROJECT_DIR$/run_main.py">
414 |       <provider selected="true" editor-type-id="text-editor">
415 |         <state relative-caret-position="301">
416 |           <caret line="220" column="20" lean-forward="true" selection-start-line="220" selection-start-column="20" selection-end-line="220" selection-end-column="20" />
417 |           <folding>
418 |             <element signature="e#0#12#0" expanded="true" />
419 |             <element signature="e#2540#5351#0" />
420 |           </folding>
421 |         </state>
422 |       </provider>
423 |     </entry>
424 |   </component>
425 | </project>


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # Variational Auto-Encoder for MNIST
  2 | - Pytorch: 0.4+
  3 | - Python: 3.6+
  4 | 
  5 | 
  6 | An Pytorch Implementation of variational auto-encoder (VAE) for MNIST descripbed in the paper:  
  7 | * [Auto-Encoding Variational Bayes](https://arxiv.org/pdf/1312.6114) by Kingma et al.
  8 | 
  9 | > This repo. is developed based on [Tensorflow-mnist-vae](https://github.com/hwalsuklee/tensorflow-mnist-VAE).
 10 | 
 11 | 
 12 | ## Results
 13 | NOTICE:
 14 | ```
 15 | tf.nn.dropout(keep_prob=0.9)
 16 | torch.nn.Dropout(p=1-keep_prob)
 17 | ```
 18 | 
 19 | ### Reproduce
 20 | Well trained VAE must be able to reproduce input image.  
 21 | Figure 5 in the paper shows reproduce performance of learned generative models for different dimensionalities.  
 22 | The following results can be reproduced with command:  
 23 | ```
 24 | python run_main.py --dim_z <each value> --num_epochs 60
 25 | ``` 
 26 | 
 27 | <table align='center'>
 28 | <tr align='center'>
 29 | <td> Input image </td>
 30 | <td> 2-D latent space </td>
 31 | <td> 5-D latent space </td>
 32 | <td> 10-D latent space </td>
 33 | <td> 20-D latent space </td>
 34 | </tr>
 35 | <tr>
 36 | <td><img src = 'res/input.jpg' height = '150px'>
 37 | <td><img src = 'res/dim_2_sample.jpg' height = '150px'>
 38 | <td><img src = 'res/dim_5_sample.jpg' height = '150px'>
 39 | <td><img src = 'res/dim_10_sample.jpg' height = '150px'>
 40 | <td><img src = 'res/dim_20_sample.jpg' height = '150px'>
 41 | </tr>
 42 | </table>
 43 | 
 44 | ### Denoising
 45 | 
 46 | When training, salt & pepper noise is added to input image, so that VAE can reduce noise and restore original input image.  
 47 | The following results can be reproduced with command:  
 48 | ```
 49 | python run_main.py --dim_z 20 --add_noise True --num_epochs 40
 50 | ```
 51 | <table align='center'>
 52 | <tr align='center'>
 53 | <td> Original input image </td>
 54 | <td> Input image with noise </td>
 55 | <td> Restored image via VAE </td>
 56 | </tr>
 57 | <tr>
 58 | <td><img src = 'res/input.jpg' height = '300px'>
 59 | <td><img src = 'res/input_noise.jpg' height = '300px'>
 60 | <td><img src = 'res/denoising.jpg' height = '300px'>
 61 | </tr>
 62 | </table>
 63 | 
 64 | ### Learned MNIST manifold
 65 | Visualizations of learned data manifold for generative models with 2-dim. latent space are given in Figure. 4 in the paper.  
 66 | The following results can be reproduced with command:  
 67 | ```
 68 | python run_main.py --dim_z 2 --num_epochs 60 --PMLR True
 69 | ```
 70 | <table align='center'>
 71 | <tr align='center'>
 72 | <td> Learned MNIST manifold </td>
 73 | <td> Distribution of labeled data  </td>
 74 | </tr>
 75 | <tr>
 76 | <td><img src = 'res/dim_2_mf.jpg' height = '400px'>
 77 | <td><img src = 'res/dim_2_cluster.jpg' height = '400px'>
 78 | </tr>
 79 | </table>
 80 | 
 81 | ## Usage
 82 | ### Prerequisites
 83 | 1. Pytorch
 84 | 2. Python packages : numpy, scipy, PIL(or Pillow), matplotlib
 85 | 
 86 | ### Command
 87 | ```
 88 | python run_main.py --dim_z <latent vector dimension>
 89 | ```
 90 | *Example*:
 91 | `python run_main.py --dim_z 20`
 92 | 
 93 | ### Arguments
 94 | *Required* :  
 95 | * `--dim_z`: Dimension of latent vector. *Default*: `20`
 96 | 
 97 | *Optional* :  
 98 | * `--results_path`: File path of output images. *Default*: `results`
 99 | * `--add_noise`: Boolean for adding salt & pepper noise to input image. *Default*: `False`
100 | * `--n_hidden`: Number of hidden units in MLP. *Default*: `500`
101 | * `--learn_rate`: Learning rate for Adam optimizer. *Default*: `1e-3`
102 | * `--num_epochs`: The number of epochs to run. *Default*: `20`
103 | * `--batch_size`: Batch size. *Default*: `128`
104 | * `--PRR`: Boolean for plot-reproduce-result. *Default*: `True`
105 | * `--PRR_n_img_x`: Number of images along x-axis. *Default*: `10`
106 | * `--PRR_n_img_y`: Number of images along y-axis. *Default*: `10`
107 | * `--PRR_resize_factor`: Resize factor for each displayed image. *Default*: `1.0`
108 | * `--PMLR`: Boolean for plot-manifold-learning-result. *Default*: `False`
109 | * `--PMLR_n_img_x`: Number of images along x-axis. *Default*: `20`
110 | * `--PMLR_n_img_y`: Number of images along y-axis. *Default*: `20`
111 | * `--PMLR_resize_factor`: Resize factor for each displayed image. *Default*: `1.0`
112 | * `--PMLR_n_samples`: Number of samples in order to get distribution of labeled data. *Default*: `5000`
113 | 
114 | ## References
115 | The implementation is based on the projects:  
116 | [1] https://github.com/oduerr/dl_tutorial/tree/master/tensorflow/vae  
117 | [2] https://github.com/fastforwardlabs/vae-tf/tree/master  
118 | [3] https://github.com/kvfrans/variational-autoencoder  
119 | [4] https://github.com/altosaar/vae
120 | 
121 | 


--------------------------------------------------------------------------------
/__pycache__/mnist_data.cpython-36.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/__pycache__/mnist_data.cpython-36.pyc


--------------------------------------------------------------------------------
/__pycache__/plot_utils.cpython-36.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/__pycache__/plot_utils.cpython-36.pyc


--------------------------------------------------------------------------------
/__pycache__/vae.cpython-36.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/__pycache__/vae.cpython-36.pyc


--------------------------------------------------------------------------------
/mnist_data.py:
--------------------------------------------------------------------------------
  1 | # Some code was borrowed from https://github.com/petewarden/tensorflow_makefile/blob/master/tensorflow/models/image/mnist/convolutional.py
  2 | 
  3 | from __future__ import absolute_import
  4 | from __future__ import division
  5 | from __future__ import print_function
  6 | 
  7 | import gzip
  8 | import os
  9 | 
 10 | import numpy
 11 | from scipy import ndimage
 12 | 
 13 | from six.moves import urllib
 14 | 
 15 | import tensorflow as tf
 16 | 
 17 | SOURCE_URL = 'http://yann.lecun.com/exdb/mnist/'
 18 | DATA_DIRECTORY = "data"
 19 | 
 20 | # Params for MNIST
 21 | IMAGE_SIZE = 28
 22 | NUM_CHANNELS = 1
 23 | PIXEL_DEPTH = 255
 24 | NUM_LABELS = 10
 25 | VALIDATION_SIZE = 5000  # Size of the validation set.
 26 | 
 27 | 
 28 | # Download MNIST data
 29 | def maybe_download(filename):
 30 |     """Download the data from Yann's website, unless it's already here."""
 31 |     if not tf.gfile.Exists(DATA_DIRECTORY):
 32 |         tf.gfile.MakeDirs(DATA_DIRECTORY)
 33 |     filepath = os.path.join(DATA_DIRECTORY, filename)
 34 |     if not tf.gfile.Exists(filepath):
 35 |         filepath, _ = urllib.request.urlretrieve(SOURCE_URL + filename, filepath)
 36 |         with tf.gfile.GFile(filepath) as f:
 37 |             size = f.size()
 38 |         print('Successfully downloaded', filename, size, 'bytes.')
 39 |     return filepath
 40 | 
 41 | 
 42 | # Extract the images
 43 | def extract_data(filename, num_images, norm_shift=False, norm_scale=True):
 44 |     """Extract the images into a 4D tensor [image index, y, x, channels].
 45 |     Values are rescaled from [0, 255] down to [-0.5, 0.5].
 46 |     """
 47 |     print('Extracting', filename)
 48 |     with gzip.open(filename) as bytestream:
 49 |         bytestream.read(16)
 50 |         buf = bytestream.read(IMAGE_SIZE * IMAGE_SIZE * num_images * NUM_CHANNELS)
 51 |         data = numpy.frombuffer(buf, dtype=numpy.uint8).astype(numpy.float32)
 52 |         if norm_shift:
 53 |             data = data - (PIXEL_DEPTH / 2.0)
 54 |         if norm_scale:
 55 |             data = data / PIXEL_DEPTH
 56 |         data = data.reshape(num_images, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS)
 57 |         data = numpy.reshape(data, [num_images, -1])
 58 |     return data
 59 | 
 60 | 
 61 | # Extract the labels
 62 | def extract_labels(filename, num_images):
 63 |     """Extract the labels into a vector of int64 label IDs."""
 64 |     print('Extracting', filename)
 65 |     with gzip.open(filename) as bytestream:
 66 |         bytestream.read(8)
 67 |         buf = bytestream.read(1 * num_images)
 68 |         labels = numpy.frombuffer(buf, dtype=numpy.uint8).astype(numpy.int64)
 69 |         num_labels_data = len(labels)
 70 |         one_hot_encoding = numpy.zeros((num_labels_data, NUM_LABELS))
 71 |         one_hot_encoding[numpy.arange(num_labels_data), labels] = 1
 72 |         one_hot_encoding = numpy.reshape(one_hot_encoding, [-1, NUM_LABELS])
 73 |     return one_hot_encoding
 74 | 
 75 | 
 76 | # Augment training data
 77 | def expend_training_data(images, labels):
 78 |     expanded_images = []
 79 |     expanded_labels = []
 80 | 
 81 |     j = 0  # counter
 82 |     for x, y in zip(images, labels):
 83 |         j = j + 1
 84 |         if j % 100 == 0:
 85 |             print('expanding data : %03d / %03d' % (j, numpy.size(images, 0)))
 86 | 
 87 |         # register original data
 88 |         expanded_images.append(x)
 89 |         expanded_labels.append(y)
 90 | 
 91 |         # get a value for the background
 92 |         # zero is the expected value, but median() is used to estimate background's value
 93 |         bg_value = numpy.median(x)  # this is regarded as background's value
 94 |         image = numpy.reshape(x, (-1, 28))
 95 | 
 96 |         for i in range(4):
 97 |             # rotate the image with random degree
 98 |             angle = numpy.random.randint(-15, 15, 1)
 99 |             new_img = ndimage.rotate(image, angle, reshape=False, cval=bg_value)
100 | 
101 |             # shift the image with random distance
102 |             shift = numpy.random.randint(-2, 2, 2)
103 |             new_img_ = ndimage.shift(new_img, shift, cval=bg_value)
104 | 
105 |             # register new training data
106 |             expanded_images.append(numpy.reshape(new_img_, 784))
107 |             expanded_labels.append(y)
108 | 
109 |     # images and labels are concatenated for random-shuffle at each epoch
110 |     # notice that pair of image and label should not be broken
111 |     expanded_train_total_data = numpy.concatenate((expanded_images, expanded_labels), axis=1)
112 |     numpy.random.shuffle(expanded_train_total_data)
113 | 
114 |     return expanded_train_total_data
115 | 
116 | 
117 | # Prepare MNISt data
118 | def prepare_MNIST_data(use_norm_shift=False, use_norm_scale=True, use_data_augmentation=False):
119 |     # Get the data.
120 |     train_data_filename = maybe_download('train-images-idx3-ubyte.gz')
121 |     train_labels_filename = maybe_download('train-labels-idx1-ubyte.gz')
122 |     test_data_filename = maybe_download('t10k-images-idx3-ubyte.gz')
123 |     test_labels_filename = maybe_download('t10k-labels-idx1-ubyte.gz')
124 | 
125 |     # Extract it into numpy arrays.
126 |     train_data = extract_data(train_data_filename, 60000, use_norm_shift, use_norm_scale)
127 |     train_labels = extract_labels(train_labels_filename, 60000)
128 |     test_data = extract_data(test_data_filename, 10000, use_norm_shift, use_norm_scale)
129 |     test_labels = extract_labels(test_labels_filename, 10000)
130 | 
131 |     # Generate a validation set.
132 |     validation_data = train_data[:VALIDATION_SIZE, :]
133 |     validation_labels = train_labels[:VALIDATION_SIZE, :]
134 |     train_data = train_data[VALIDATION_SIZE:, :]
135 |     train_labels = train_labels[VALIDATION_SIZE:, :]
136 | 
137 |     # Concatenate train_data & train_labels for random shuffle
138 |     if use_data_augmentation:
139 |         train_total_data = expend_training_data(train_data, train_labels)
140 |     else:
141 |         train_total_data = numpy.concatenate((train_data, train_labels), axis=1)
142 | 
143 |     train_size = train_total_data.shape[0]
144 | 
145 |     return train_total_data, train_size, validation_data, validation_labels, test_data, test_labels
146 | 


--------------------------------------------------------------------------------
/plot_utils.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | import matplotlib.pyplot as plt
  3 | from scipy.misc import imsave
  4 | from scipy.misc import imresize
  5 | 
  6 | 
  7 | class Plot_Reproduce_Performance():
  8 |     def __init__(self, DIR, n_img_x=8, n_img_y=8, img_w=28, img_h=28, resize_factor=1.0):
  9 |         self.DIR = DIR
 10 | 
 11 |         assert n_img_x > 0 and n_img_y > 0
 12 | 
 13 |         self.n_img_x = n_img_x
 14 |         self.n_img_y = n_img_y
 15 |         self.n_tot_imgs = n_img_x * n_img_y
 16 | 
 17 |         assert img_w > 0 and img_h > 0
 18 | 
 19 |         self.img_w = img_w
 20 |         self.img_h = img_h
 21 | 
 22 |         assert resize_factor > 0
 23 | 
 24 |         self.resize_factor = resize_factor
 25 | 
 26 |     def save_images(self, images, name='result.jpg'):
 27 |         images = images.reshape(self.n_img_x * self.n_img_y, self.img_h, self.img_w)
 28 |         imsave(self.DIR + "/" + name, self._merge(images, [self.n_img_y, self.n_img_x]))
 29 | 
 30 |     def _merge(self, images, size):
 31 |         h, w = images.shape[1], images.shape[2]
 32 | 
 33 |         h_ = int(h * self.resize_factor)
 34 |         w_ = int(w * self.resize_factor)
 35 | 
 36 |         img = np.zeros((h_ * size[0], w_ * size[1]))
 37 | 
 38 |         for idx, image in enumerate(images):
 39 |             i = int(idx % size[1])
 40 |             j = int(idx / size[1])
 41 | 
 42 |             image_ = imresize(image, size=(w_, h_), interp='bicubic')
 43 | 
 44 |             img[j * h_:j * h_ + h_, i * w_:i * w_ + w_] = image_
 45 | 
 46 |         return img
 47 | 
 48 | 
 49 | class Plot_Manifold_Learning_Result():
 50 |     def __init__(self, DIR, n_img_x=20, n_img_y=20, img_w=28, img_h=28, resize_factor=1.0, z_range=4):
 51 |         self.DIR = DIR
 52 | 
 53 |         assert n_img_x > 0 and n_img_y > 0
 54 | 
 55 |         self.n_img_x = n_img_x
 56 |         self.n_img_y = n_img_y
 57 |         self.n_tot_imgs = n_img_x * n_img_y
 58 | 
 59 |         assert img_w > 0 and img_h > 0
 60 | 
 61 |         self.img_w = img_w
 62 |         self.img_h = img_h
 63 | 
 64 |         assert resize_factor > 0
 65 | 
 66 |         self.resize_factor = resize_factor
 67 | 
 68 |         assert z_range > 0
 69 |         self.z_range = z_range
 70 | 
 71 |         self._set_latent_vectors()
 72 | 
 73 |     def _set_latent_vectors(self):
 74 |         # z1 = np.linspace(-self.z_range, self.z_range, self.n_img_y)
 75 |         # z2 = np.linspace(-self.z_range, self.z_range, self.n_img_x)
 76 |         #
 77 |         # z = np.array(np.meshgrid(z1, z2))
 78 |         # z = z.reshape([-1, 2])
 79 | 
 80 |         # borrowed from https://github.com/fastforwardlabs/vae-tf/blob/master/plot.py
 81 |         z = np.rollaxis(
 82 |             np.mgrid[self.z_range:-self.z_range:self.n_img_y * 1j, self.z_range:-self.z_range:self.n_img_x * 1j], 0, 3)
 83 |         # z1 = np.rollaxis(np.mgrid[1:-1:self.n_img_y * 1j, 1:-1:self.n_img_x * 1j], 0, 3)
 84 |         # z = z1**2
 85 |         # z[z1<0] *= -1
 86 |         #
 87 |         # z = z*self.z_range
 88 | 
 89 |         self.z = z.reshape([-1, 2])
 90 | 
 91 |     def save_images(self, images, name='result.jpg'):
 92 |         images = images.reshape(self.n_img_x * self.n_img_y, self.img_h, self.img_w)
 93 |         imsave(self.DIR + "/" + name, self._merge(images, [self.n_img_y, self.n_img_x]))
 94 | 
 95 |     def _merge(self, images, size):
 96 |         h, w = images.shape[1], images.shape[2]
 97 | 
 98 |         h_ = int(h * self.resize_factor)
 99 |         w_ = int(w * self.resize_factor)
100 | 
101 |         img = np.zeros((h_ * size[0], w_ * size[1]))
102 | 
103 |         for idx, image in enumerate(images):
104 |             i = int(idx % size[1])
105 |             j = int(idx / size[1])
106 | 
107 |             image_ = imresize(image, size=(w_, h_), interp='bicubic')
108 | 
109 |             img[j * h_:j * h_ + h_, i * w_:i * w_ + w_] = image_
110 | 
111 |         return img
112 | 
113 |     # borrowed from https://github.com/ykwon0407/variational_autoencoder/blob/master/variational_bayes.ipynb
114 |     def save_scattered_image(self, z, id, name='scattered_image.jpg'):
115 |         N = 10
116 |         plt.figure(figsize=(8, 6))
117 |         plt.scatter(z[:, 0], z[:, 1], c=np.argmax(id, 1), marker='o', edgecolor='none', cmap=discrete_cmap(N, 'jet'))
118 |         plt.colorbar(ticks=range(N))
119 |         axes = plt.gca()
120 |         axes.set_xlim([-self.z_range - 2, self.z_range + 2])
121 |         axes.set_ylim([-self.z_range - 2, self.z_range + 2])
122 |         plt.grid(True)
123 |         plt.savefig(self.DIR + "/" + name)
124 | 
125 | 
126 | # borrowed from https://gist.github.com/jakevdp/91077b0cae40f8f8244a
127 | def discrete_cmap(N, base_cmap=None):
128 |     """Create an N-bin discrete colormap from the specified input map"""
129 | 
130 |     # Note that if base_cmap is a string or None, you can simply do
131 |     #    return plt.cm.get_cmap(base_cmap, N)
132 |     # The following works for string, None, or a colormap instance:
133 | 
134 |     base = plt.cm.get_cmap(base_cmap)
135 |     color_list = base(np.linspace(0, 1, N))
136 |     cmap_name = base.name + str(N)
137 |     return base.from_list(cmap_name, color_list, N)
138 | 


--------------------------------------------------------------------------------
/res/denoising.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/res/denoising.jpg


--------------------------------------------------------------------------------
/res/dim_10_sample.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/res/dim_10_sample.jpg


--------------------------------------------------------------------------------
/res/dim_20_sample.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/res/dim_20_sample.jpg


--------------------------------------------------------------------------------
/res/dim_2_cluster.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/res/dim_2_cluster.jpg


--------------------------------------------------------------------------------
/res/dim_2_mf.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/res/dim_2_mf.jpg


--------------------------------------------------------------------------------
/res/dim_2_sample.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/res/dim_2_sample.jpg


--------------------------------------------------------------------------------
/res/dim_5_sample.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/res/dim_5_sample.jpg


--------------------------------------------------------------------------------
/res/input.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/res/input.jpg


--------------------------------------------------------------------------------
/res/input_noise.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/dragen1860/pytorch-mnist-vae/3ebd2acac8d0d6cc68d0f9b4d04926b682d6dc44/res/input_noise.jpg


--------------------------------------------------------------------------------
/run_main.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import numpy as np
  3 | import mnist_data
  4 | import os
  5 | import vae
  6 | import plot_utils
  7 | import glob
  8 | 
  9 | import argparse
 10 | 
 11 | IMAGE_SIZE_MNIST = 28
 12 | 
 13 | """parsing and configuration"""
 14 | 
 15 | 
 16 | def parse_args():
 17 |     desc = "Pytorch implementation of 'Variational AutoEncoder (VAE)'"
 18 |     parser = argparse.ArgumentParser(description=desc)
 19 | 
 20 |     parser.add_argument('--results_path', type=str, default='results',
 21 |                         help='File path of output images')
 22 | 
 23 |     parser.add_argument('--add_noise', type=bool, default=False,
 24 |                         help='Boolean for adding salt & pepper noise to input image')
 25 | 
 26 |     parser.add_argument('--dim_z', type=int, default='20', help='Dimension of latent vector', required=True)
 27 | 
 28 |     parser.add_argument('--n_hidden', type=int, default=500, help='Number of hidden units in MLP')
 29 | 
 30 |     parser.add_argument('--learn_rate', type=float, default=1e-3, help='Learning rate for Adam optimizer')
 31 | 
 32 |     parser.add_argument('--num_epochs', type=int, default=20, help='The number of epochs to run')
 33 | 
 34 |     parser.add_argument('--batch_size', type=int, default=128, help='Batch size')
 35 | 
 36 |     parser.add_argument('--PRR', type=bool, default=True,
 37 |                         help='Boolean for plot-reproduce-result')
 38 | 
 39 |     parser.add_argument('--PRR_n_img_x', type=int, default=10,
 40 |                         help='Number of images along x-axis')
 41 | 
 42 |     parser.add_argument('--PRR_n_img_y', type=int, default=10,
 43 |                         help='Number of images along y-axis')
 44 | 
 45 |     parser.add_argument('--PRR_resize_factor', type=float, default=1.0,
 46 |                         help='Resize factor for each displayed image')
 47 | 
 48 |     parser.add_argument('--PMLR', type=bool, default=False,
 49 |                         help='Boolean for plot-manifold-learning-result')
 50 | 
 51 |     parser.add_argument('--PMLR_n_img_x', type=int, default=20,
 52 |                         help='Number of images along x-axis')
 53 | 
 54 |     parser.add_argument('--PMLR_n_img_y', type=int, default=20,
 55 |                         help='Number of images along y-axis')
 56 | 
 57 |     parser.add_argument('--PMLR_resize_factor', type=float, default=1.0,
 58 |                         help='Resize factor for each displayed image')
 59 | 
 60 |     parser.add_argument('--PMLR_z_range', type=float, default=2.0,
 61 |                         help='Range for unifomly distributed latent vector')
 62 | 
 63 |     parser.add_argument('--PMLR_n_samples', type=int, default=5000,
 64 |                         help='Number of samples in order to get distribution of labeled data')
 65 | 
 66 |     return check_args(parser.parse_args())
 67 | 
 68 | 
 69 | 
 70 | def check_args(args):
 71 |     # --results_path
 72 |     try:
 73 |         os.mkdir(args.results_path)
 74 |     except(FileExistsError):
 75 |         pass
 76 |     # delete all existing files
 77 |     files = glob.glob(args.results_path + '/*')
 78 |     for f in files:
 79 |         os.remove(f)
 80 | 
 81 |     # --add_noise
 82 |     try:
 83 |         assert args.add_noise == True or args.add_noise == False
 84 |     except:
 85 |         print('add_noise must be boolean type')
 86 |         return None
 87 | 
 88 |     # --dim-z
 89 |     try:
 90 |         assert args.dim_z > 0
 91 |     except:
 92 |         print('dim_z must be positive integer')
 93 |         return None
 94 | 
 95 |     # --n_hidden
 96 |     try:
 97 |         assert args.n_hidden >= 1
 98 |     except:
 99 |         print('number of hidden units must be larger than one')
100 | 
101 |     # --learn_rate
102 |     try:
103 |         assert args.learn_rate > 0
104 |     except:
105 |         print('learning rate must be positive')
106 | 
107 |     # --num_epochs
108 |     try:
109 |         assert args.num_epochs >= 1
110 |     except:
111 |         print('number of epochs must be larger than or equal to one')
112 | 
113 |     # --batch_size
114 |     try:
115 |         assert args.batch_size >= 1
116 |     except:
117 |         print('batch size must be larger than or equal to one')
118 | 
119 |     # --PRR
120 |     try:
121 |         assert args.PRR == True or args.PRR == False
122 |     except:
123 |         print('PRR must be boolean type')
124 |         return None
125 | 
126 |     if args.PRR == True:
127 |         # --PRR_n_img_x, --PRR_n_img_y
128 |         try:
129 |             assert args.PRR_n_img_x >= 1 and args.PRR_n_img_y >= 1
130 |         except:
131 |             print('PRR : number of images along each axis must be larger than or equal to one')
132 | 
133 |         # --PRR_resize_factor
134 |         try:
135 |             assert args.PRR_resize_factor > 0
136 |         except:
137 |             print('PRR : resize factor for each displayed image must be positive')
138 | 
139 |     # --PMLR
140 |     try:
141 |         assert args.PMLR == True or args.PMLR == False
142 |     except:
143 |         print('PMLR must be boolean type')
144 |         return None
145 | 
146 |     if args.PMLR == True:
147 |         try:
148 |             assert args.dim_z == 2
149 |         except:
150 |             print('PMLR : dim_z must be two')
151 | 
152 |         # --PMLR_n_img_x, --PMLR_n_img_y
153 |         try:
154 |             assert args.PMLR_n_img_x >= 1 and args.PMLR_n_img_y >= 1
155 |         except:
156 |             print('PMLR : number of images along each axis must be larger than or equal to one')
157 | 
158 |         # --PMLR_resize_factor
159 |         try:
160 |             assert args.PMLR_resize_factor > 0
161 |         except:
162 |             print('PMLR : resize factor for each displayed image must be positive')
163 | 
164 |         # --PMLR_z_range
165 |         try:
166 |             assert args.PMLR_z_range > 0
167 |         except:
168 |             print('PMLR : range for unifomly distributed latent vector must be positive')
169 | 
170 |         # --PMLR_n_samples
171 |         try:
172 |             assert args.PMLR_n_samples > 100
173 |         except:
174 |             print('PMLR : Number of samples in order to get distribution of labeled data must be large enough')
175 | 
176 |     return args
177 | 
178 | 
179 | 
180 | 
181 | 
182 | def main(args):
183 | 
184 | 
185 |     # torch.manual_seed(222)
186 |     # torch.cuda.manual_seed_all(222)
187 |     # np.random.seed(222)
188 | 
189 | 
190 |     device = torch.device('cuda')
191 | 
192 |     RESULTS_DIR = args.results_path
193 |     ADD_NOISE = args.add_noise
194 |     n_hidden = args.n_hidden
195 |     dim_img = IMAGE_SIZE_MNIST ** 2  # number of pixels for a MNIST image
196 |     dim_z = args.dim_z
197 | 
198 |     # train
199 |     n_epochs = args.num_epochs
200 |     batch_size = args.batch_size
201 |     learn_rate = args.learn_rate
202 | 
203 |     # Plot
204 |     PRR = args.PRR  # Plot Reproduce Result
205 |     PRR_n_img_x = args.PRR_n_img_x  # number of images along x-axis in a canvas
206 |     PRR_n_img_y = args.PRR_n_img_y  # number of images along y-axis in a canvas
207 |     PRR_resize_factor = args.PRR_resize_factor  # resize factor for each image in a canvas
208 | 
209 |     PMLR = args.PMLR  # Plot Manifold Learning Result
210 |     PMLR_n_img_x = args.PMLR_n_img_x  # number of images along x-axis in a canvas
211 |     PMLR_n_img_y = args.PMLR_n_img_y  # number of images along y-axis in a canvas
212 |     PMLR_resize_factor = args.PMLR_resize_factor  # resize factor for each image in a canvas
213 |     PMLR_z_range = args.PMLR_z_range  # range for random latent vector
214 |     PMLR_n_samples = args.PMLR_n_samples  # number of labeled samples to plot a map from input data space to the latent space
215 | 
216 |     """ prepare MNIST data """
217 |     train_total_data, train_size, _, _, test_data, test_labels = mnist_data.prepare_MNIST_data()
218 |     n_samples = train_size
219 | 
220 |     """ create network """
221 |     keep_prob = 0.99
222 |     encoder = vae.Encoder(dim_img, n_hidden, dim_z, keep_prob).to(device)
223 |     decoder = vae.Decoder(dim_z, n_hidden, dim_img, keep_prob).to(device)
224 |     # + operator will return but .extend is inplace no return.
225 |     optimizer = torch.optim.Adam(list(encoder.parameters()) + list(decoder.parameters()), lr=learn_rate)
226 |     # vae.init_weights(encoder, decoder)
227 | 
228 |     """ training """
229 |     # Plot for reproduce performance
230 |     if PRR:
231 |         PRR = plot_utils.Plot_Reproduce_Performance(RESULTS_DIR, PRR_n_img_x, PRR_n_img_y, IMAGE_SIZE_MNIST,
232 |                                                     IMAGE_SIZE_MNIST, PRR_resize_factor)
233 | 
234 |         x_PRR = test_data[0:PRR.n_tot_imgs, :]
235 | 
236 |         x_PRR_img = x_PRR.reshape(PRR.n_tot_imgs, IMAGE_SIZE_MNIST, IMAGE_SIZE_MNIST)
237 |         PRR.save_images(x_PRR_img, name='input.jpg')
238 |         print('saved:', 'input.jpg')
239 | 
240 |         if ADD_NOISE:
241 |             x_PRR = x_PRR * np.random.randint(2, size=x_PRR.shape)
242 |             x_PRR += np.random.randint(2, size=x_PRR.shape)
243 | 
244 |             x_PRR_img = x_PRR.reshape(PRR.n_tot_imgs, IMAGE_SIZE_MNIST, IMAGE_SIZE_MNIST)
245 |             PRR.save_images(x_PRR_img, name='input_noise.jpg')
246 |             print('saved:', 'input_noise.jpg')
247 | 
248 |         x_PRR = torch.from_numpy(x_PRR).float().to(device)
249 | 
250 |     # Plot for manifold learning result
251 |     if PMLR and dim_z == 2:
252 | 
253 |         PMLR = plot_utils.Plot_Manifold_Learning_Result(RESULTS_DIR, PMLR_n_img_x, PMLR_n_img_y, IMAGE_SIZE_MNIST,
254 |                                                         IMAGE_SIZE_MNIST, PMLR_resize_factor, PMLR_z_range)
255 | 
256 |         x_PMLR = test_data[0:PMLR_n_samples, :]
257 |         id_PMLR = test_labels[0:PMLR_n_samples, :]
258 | 
259 |         if ADD_NOISE:
260 |             x_PMLR = x_PMLR * np.random.randint(2, size=x_PMLR.shape)
261 |             x_PMLR += np.random.randint(2, size=x_PMLR.shape)
262 | 
263 | 
264 |         z_ = torch.from_numpy(PMLR.z).float().to(device)
265 |         x_PMLR = torch.from_numpy(x_PMLR).float().to(device)
266 | 
267 | 
268 |     # train
269 |     total_batch = int(n_samples / batch_size)
270 |     min_tot_loss = np.inf
271 |     for epoch in range(n_epochs):
272 | 
273 |         # Random shuffling
274 |         np.random.shuffle(train_total_data)
275 |         train_data_ = train_total_data[:, :-mnist_data.NUM_LABELS]
276 | 
277 |         # Loop over all batches
278 |         encoder.train()
279 |         decoder.train()
280 |         for i in range(total_batch):
281 |             # Compute the offset of the current minibatch in the data.
282 |             offset = (i * batch_size) % (n_samples)
283 |             batch_xs_input = train_data_[offset:(offset + batch_size), :]
284 | 
285 |             batch_xs_target = batch_xs_input
286 | 
287 |             # add salt & pepper noise
288 |             if ADD_NOISE:
289 |                 batch_xs_input = batch_xs_input * np.random.randint(2, size=batch_xs_input.shape)
290 |                 batch_xs_input += np.random.randint(2, size=batch_xs_input.shape)
291 | 
292 |             batch_xs_input, batch_xs_target = torch.from_numpy(batch_xs_input).float().to(device), \
293 |                                               torch.from_numpy(batch_xs_target).float().to(device)
294 | 
295 |             assert not torch.isnan(batch_xs_input).any()
296 |             assert not torch.isnan(batch_xs_target).any()
297 | 
298 |             y, z, tot_loss, loss_likelihood, loss_divergence = \
299 |                                         vae.get_loss(encoder, decoder, batch_xs_input, batch_xs_target)
300 | 
301 |             optimizer.zero_grad()
302 |             tot_loss.backward()
303 |             optimizer.step()
304 | 
305 | 
306 | 
307 |             # print cost every epoch
308 |         print("epoch %d: L_tot %03.2f L_likelihood %03.2f L_divergence %03.2f" % (
309 |                                                 epoch, tot_loss.item(), loss_likelihood.item(), loss_divergence.item()))
310 | 
311 | 
312 | 
313 | 
314 |         encoder.eval()
315 |         decoder.eval()
316 |         # if minimum loss is updated or final epoch, plot results
317 |         if min_tot_loss > tot_loss.item() or epoch + 1 == n_epochs:
318 |             min_tot_loss = tot_loss.item()
319 | 
320 |             # Plot for reproduce performance
321 |             if PRR:
322 |                 y_PRR = vae.get_ae(encoder, decoder, x_PRR)
323 | 
324 |                 y_PRR_img = y_PRR.reshape(PRR.n_tot_imgs, IMAGE_SIZE_MNIST, IMAGE_SIZE_MNIST)
325 |                 PRR.save_images(y_PRR_img.detach().cpu().numpy(), name="/PRR_epoch_%02d" % (epoch) + ".jpg")
326 |                 print('saved:', "/PRR_epoch_%02d" % (epoch) + ".jpg")
327 | 
328 |             # Plot for manifold learning result
329 |             if PMLR and dim_z == 2:
330 |                 y_PMLR = decoder(z_)
331 | 
332 |                 y_PMLR_img = y_PMLR.reshape(PMLR.n_tot_imgs, IMAGE_SIZE_MNIST, IMAGE_SIZE_MNIST)
333 |                 PMLR.save_images(y_PMLR_img.detach().cpu().numpy(), name="/PMLR_epoch_%02d" % (epoch) + ".jpg")
334 |                 print('saved:', "/PMLR_epoch_%02d" % (epoch) + ".jpg")
335 | 
336 |                 # plot distribution of labeled images
337 |                 z_PMLR = vae.get_z(encoder, x_PMLR)
338 |                 PMLR.save_scattered_image(z_PMLR.detach().cpu().numpy(), id_PMLR,
339 |                                           name="/PMLR_map_epoch_%02d" % (epoch) + ".jpg")
340 |                 print('saved:', "/PMLR_map_epoch_%02d" % (epoch) + ".jpg")
341 | 
342 | 
343 | if __name__ == '__main__':
344 | 
345 |     # parse arguments
346 |     args = parse_args()
347 |     if args is None:
348 |         exit()
349 | 
350 |     # main
351 |     main(args)
352 | 


--------------------------------------------------------------------------------
/test.py:
--------------------------------------------------------------------------------
 1 | import  torch
 2 | import  tensorflow as tf
 3 | from    torch.nn import functional as F
 4 | 
 5 | 
 6 | tf.enable_eager_execution()
 7 | 
 8 | 
 9 | 
10 | 
11 | 
12 | 
13 | 
14 | 
15 | def main():
16 | 
17 |     y = torch.rand((3,20))
18 |     x = torch.randint(0, 2, (3,20))
19 | 
20 |     print(x)
21 |     print(y)
22 | 
23 |     loss = torch.sum(x * torch.log(y) + (1 - x) * torch.log(1 - y)) / 3
24 |     loss2 = -F.binary_cross_entropy(y, x, reduction='sum') / 3
25 | 
26 | 
27 |     x = tf.Variable(x.numpy())
28 |     y = tf.Variable(y.numpy())
29 | 
30 |     loss3 = tf.reduce_sum(x * tf.log(y) + (1 - x) * tf.log(1 - y)) / 3
31 |     y = tf.clip_by_value(y, 1e-8, 1 - 1e-8)
32 |     loss4 = tf.reduce_sum(x * tf.log(y) + (1 - x) * tf.log(1 - y)) / 3
33 | 
34 | 
35 |     print(loss, loss2, loss3, loss4)
36 | 
37 | 
38 | if __name__ == '__main__':
39 |     main()
40 | 
41 | 
42 | 


--------------------------------------------------------------------------------
/vae.py:
--------------------------------------------------------------------------------
  1 | import  torch
  2 | from    torch import nn
  3 | from    torch.nn import functional as F
  4 | 
  5 | 
  6 | class Encoder(nn.Module):
  7 | 
  8 | 
  9 |     def __init__(self, imgsz, n_hidden, n_output, keep_prob):
 10 |         super(Encoder, self).__init__()
 11 | 
 12 |         self.imgsz = imgsz
 13 |         self.n_hidden = n_hidden
 14 |         self.n_output = n_output
 15 |         self.keep_prob = keep_prob
 16 | 
 17 |         self.net = nn.Sequential(
 18 |             nn.Linear(imgsz, n_hidden),
 19 |             nn.ELU(inplace=True),
 20 |             nn.Dropout(1-keep_prob),
 21 | 
 22 |             nn.Linear(n_hidden, n_hidden),
 23 |             nn.Tanh(),
 24 |             nn.Dropout(1-keep_prob),
 25 | 
 26 |             nn.Linear(n_hidden, n_output*2)
 27 | 
 28 |         )
 29 | 
 30 |     def forward(self, x):
 31 |         """
 32 | 
 33 |         :param x:
 34 |         :return:
 35 |         """
 36 |         mu_sigma = self.net(x)
 37 | 
 38 | 
 39 |         # The mean parameter is unconstrained
 40 |         mean = mu_sigma[:, :self.n_output]
 41 |         # The standard deviation must be positive. Parametrize with a softplus and
 42 |         # add a small epsilon for numerical stability
 43 |         stddev = 1e-6 + F.softplus(mu_sigma[:, self.n_output:])
 44 | 
 45 | 
 46 |         return mean, stddev
 47 | 
 48 | 
 49 | 
 50 | class Decoder(nn.Module):
 51 | 
 52 | 
 53 |     def __init__(self, dim_z, n_hidden, n_output, keep_prob):
 54 |         super(Decoder, self).__init__()
 55 | 
 56 |         self.dim_z = dim_z
 57 |         self.n_hidden = n_hidden
 58 |         self.n_output = n_output
 59 |         self.keep_prob = keep_prob
 60 | 
 61 |         self.net = nn.Sequential(
 62 |             nn.Linear(dim_z, n_hidden),
 63 |             nn.Tanh(),
 64 |             nn.Dropout(1-keep_prob),
 65 | 
 66 |             nn.Linear(n_hidden, n_hidden),
 67 |             nn.ELU(),
 68 |             nn.Dropout(1-keep_prob),
 69 | 
 70 |             nn.Linear(n_hidden, n_output),
 71 |             nn.Sigmoid()
 72 |         )
 73 | 
 74 |     def forward(self, h):
 75 |         """
 76 | 
 77 |         :param h:
 78 |         :return:
 79 |         """
 80 |         return self.net(h)
 81 | 
 82 | 
 83 | 
 84 | def init_weights(encoder, decoder):
 85 | 
 86 |     def init_(m):
 87 |         with torch.no_grad():
 88 |             if isinstance(m, nn.Linear):
 89 |                 nn.init.kaiming_normal_(m.weight)
 90 |                 nn.init.constant_(m.bias, 0.0)
 91 | 
 92 |     for m in encoder.modules():
 93 |         m.apply(init_)
 94 |     for m in decoder.modules():
 95 |         m.apply(init_)
 96 | 
 97 |     print('weights inited!')
 98 | 
 99 | 
100 | 
101 | def get_ae(encoder, decoder, x):
102 |     # encoding
103 |     mu, sigma = encoder(x)
104 |     # sampling by re-parameterization technique
105 |     z = mu + sigma * torch.randn_like(mu)
106 | 
107 |     # decoding
108 |     y = decoder(z)
109 |     y = torch.clamp(y, 1e-8, 1 - 1e-8)
110 | 
111 |     return y
112 | 
113 | 
114 | 
115 | def get_z(encoder, x):
116 | 
117 |     # encoding
118 |     mu, sigma = encoder(x)
119 |     # sampling by re-parameterization technique
120 |     z = mu + sigma * torch.randn_like(mu)
121 | 
122 |     return z
123 | 
124 | def get_loss(encoder, decoder, x, x_target):
125 |     """
126 | 
127 |     :param encoder:
128 |     :param decoder:
129 |     :param x: input
130 |     :param x_hat: target
131 |     :param dim_img:
132 |     :param dim_z:
133 |     :param n_hidden:
134 |     :param keep_prob:
135 |     :return:
136 |     """
137 |     batchsz = x.size(0)
138 |     # encoding
139 |     mu, sigma = encoder(x)
140 |     # sampling by re-parameterization technique
141 |     z = mu + sigma * torch.randn_like(mu)
142 | 
143 |     # decoding
144 |     y = decoder(z)
145 |     y = torch.clamp(y, 1e-8, 1 - 1e-8)
146 | 
147 | 
148 |     # loss
149 |     # marginal_likelihood2 = torch.sum(x_target * torch.log(y) + (1 - x_target) * torch.log(1 - y)) / batchsz
150 |     marginal_likelihood = -F.binary_cross_entropy(y, x_target, reduction='sum') / batchsz
151 |     # print(marginal_likelihood2.item(), marginal_likelihood.item())
152 | 
153 |     KL_divergence = 0.5 * torch.sum(
154 |                                 torch.pow(mu, 2) +
155 |                                 torch.pow(sigma, 2) -
156 |                                 torch.log(1e-8 + torch.pow(sigma, 2)) - 1
157 |                                ).sum() / batchsz
158 | 
159 |     ELBO = marginal_likelihood - KL_divergence
160 | 
161 |     loss = -ELBO
162 | 
163 |     return y, z, loss, marginal_likelihood, KL_divergence
164 | 
165 | 
166 | # # Gaussian MLP as encoder
167 | # def gaussian_MLP_encoder(x, n_hidden, n_output, keep_prob):
168 | #     with tf.variable_scope("gaussian_MLP_encoder"):
169 | #         # initializers
170 | #         w_init = tf.contrib.layers.variance_scaling_initializer()
171 | #         b_init = tf.constant_initializer(0.)
172 | #
173 | #         # 1st hidden layer
174 | #         w0 = tf.get_variable('w0', [x.get_shape()[1], n_hidden], initializer=w_init)
175 | #         b0 = tf.get_variable('b0', [n_hidden], initializer=b_init)
176 | #         h0 = tf.matmul(x, w0) + b0
177 | #         h0 = tf.nn.elu(h0)
178 | #         h0 = tf.nn.dropout(h0, keep_prob)
179 | #
180 | #         # 2nd hidden layer
181 | #         w1 = tf.get_variable('w1', [h0.get_shape()[1], n_hidden], initializer=w_init)
182 | #         b1 = tf.get_variable('b1', [n_hidden], initializer=b_init)
183 | #         h1 = tf.matmul(h0, w1) + b1
184 | #         h1 = tf.nn.tanh(h1)
185 | #         h1 = tf.nn.dropout(h1, keep_prob)
186 | #
187 | #         # output layer
188 | #         # borrowed from https: // github.com / altosaar / vae / blob / master / vae.py
189 | #         wo = tf.get_variable('wo', [h1.get_shape()[1], n_output * 2], initializer=w_init)
190 | #         bo = tf.get_variable('bo', [n_output * 2], initializer=b_init)
191 | #         gaussian_params = tf.matmul(h1, wo) + bo
192 | #
193 | #         # The mean parameter is unconstrained
194 | #         mean = gaussian_params[:, :n_output]
195 | #         # The standard deviation must be positive. Parametrize with a softplus and
196 | #         # add a small epsilon for numerical stability
197 | #         stddev = 1e-6 + tf.nn.softplus(gaussian_params[:, n_output:])
198 | #
199 | #     return mean, stddev
200 | #
201 | #
202 | # # Bernoulli MLP as decoder
203 | # def bernoulli_MLP_decoder(z, n_hidden, n_output, keep_prob, reuse=False):
204 | #     with tf.variable_scope("bernoulli_MLP_decoder", reuse=reuse):
205 | #         # initializers
206 | #         w_init = tf.contrib.layers.variance_scaling_initializer()
207 | #         b_init = tf.constant_initializer(0.)
208 | #
209 | #         # 1st hidden layer
210 | #         w0 = tf.get_variable('w0', [z.get_shape()[1], n_hidden], initializer=w_init)
211 | #         b0 = tf.get_variable('b0', [n_hidden], initializer=b_init)
212 | #         h0 = tf.matmul(z, w0) + b0
213 | #         h0 = tf.nn.tanh(h0)
214 | #         h0 = tf.nn.dropout(h0, keep_prob)
215 | #
216 | #         # 2nd hidden layer
217 | #         w1 = tf.get_variable('w1', [h0.get_shape()[1], n_hidden], initializer=w_init)
218 | #         b1 = tf.get_variable('b1', [n_hidden], initializer=b_init)
219 | #         h1 = tf.matmul(h0, w1) + b1
220 | #         h1 = tf.nn.elu(h1)
221 | #         h1 = tf.nn.dropout(h1, keep_prob)
222 | #
223 | #         # output layer-mean
224 | #         wo = tf.get_variable('wo', [h1.get_shape()[1], n_output], initializer=w_init)
225 | #         bo = tf.get_variable('bo', [n_output], initializer=b_init)
226 | #         y = tf.sigmoid(tf.matmul(h1, wo) + bo)
227 | #
228 | #     return y
229 | #
230 | #
231 | # # Gateway
232 | # def autoencoder2(x_hat, x, dim_img, dim_z, n_hidden, keep_prob):
233 | #     # encoding
234 | #     mu, sigma = gaussian_MLP_encoder(x_hat, n_hidden, dim_z, keep_prob)
235 | #
236 | #     # sampling by re-parameterization technique
237 | #     z = mu + sigma * tf.random_normal(tf.shape(mu), 0, 1, dtype=tf.float32)
238 | #
239 | #     # decoding
240 | #     y = bernoulli_MLP_decoder(z, n_hidden, dim_img, keep_prob)
241 | #     y = tf.clip_by_value(y, 1e-8, 1 - 1e-8)
242 | #
243 | #     # loss
244 | #     marginal_likelihood = tf.reduce_sum(x * tf.log(y) + (1 - x) * tf.log(1 - y), 1)
245 | #     KL_divergence = 0.5 * tf.reduce_sum(tf.square(mu) + tf.square(sigma) - tf.log(1e-8 + tf.square(sigma)) - 1, 1)
246 | #
247 | #     marginal_likelihood = tf.reduce_mean(marginal_likelihood)
248 | #     KL_divergence = tf.reduce_mean(KL_divergence)
249 | #
250 | #     ELBO = marginal_likelihood - KL_divergence
251 | #
252 | #     loss = -ELBO
253 | #
254 | #     return y, z, loss, -marginal_likelihood, KL_divergence
255 | #
256 | #
257 | # def decoder(z, dim_img, n_hidden):
258 | #     y = bernoulli_MLP_decoder(z, n_hidden, dim_img, 1.0, reuse=True)
259 | #
260 | #     return y
261 | 


--------------------------------------------------------------------------------
/vae2.py:
--------------------------------------------------------------------------------
  1 | import  torch
  2 | from    torch import nn
  3 | from    torch.nn import functional as F
  4 | 
  5 | 
  6 | 
  7 | 
  8 | class Encoder(torch.nn.Module):
  9 |     def __init__(self, D_in, H, D_out, keep_prob=0):
 10 |         super(Encoder, self).__init__()
 11 | 
 12 |         self.linear1 = torch.nn.Linear(D_in, H)
 13 |         self.linear2 = torch.nn.Linear(H, 32)
 14 | 
 15 |         self._enc_mu = torch.nn.Linear(32, D_out)
 16 |         self._enc_log_sigma = torch.nn.Linear(32, D_out)
 17 | 
 18 |     def forward(self, x):
 19 |         x = F.relu(self.linear1(x))
 20 |         x = F.relu(self.linear2(x))
 21 | 
 22 |         # print(x.shape)
 23 | 
 24 |         return self._enc_mu(x), self._enc_log_sigma(x)
 25 | 
 26 | 
 27 | 
 28 | 
 29 | class Decoder(torch.nn.Module):
 30 |     def __init__(self, D_in, H, D_out, keep_prob=0):
 31 |         super(Decoder, self).__init__()
 32 |         self.linear1 = torch.nn.Linear(D_in, H)
 33 |         self.linear2 = torch.nn.Linear(H, D_out)
 34 | 
 35 |     def forward(self, x):
 36 |         x = F.relu(self.linear1(x))
 37 |         return F.relu(self.linear2(x))
 38 | 
 39 | 
 40 | def get_ae(encoder, decoder, x):
 41 |     # encoding
 42 |     mu, log_sigma = encoder(x)
 43 |     sigma = torch.exp(log_sigma)
 44 |     # sampling by re-parameterization technique
 45 |     z = mu + sigma * torch.randn_like(mu)
 46 | 
 47 |     # decoding
 48 |     y = decoder(z)
 49 |     # y = torch.clamp(y, 1e-8, 1 - 1e-8)
 50 | 
 51 |     return y
 52 | 
 53 | 
 54 | 
 55 | def get_z(encoder, x):
 56 | 
 57 |     # encoding
 58 |     mu, log_sigma = encoder(x)
 59 |     sigma = torch.exp(log_sigma)
 60 |     # sampling by re-parameterization technique
 61 |     z = mu + sigma * torch.randn_like(mu)
 62 | 
 63 |     return z
 64 | 
 65 | 
 66 | 
 67 | def get_loss(encoder, decoder, x, x_target):
 68 |     """
 69 | 
 70 |     :param encoder:
 71 |     :param decoder:
 72 |     :param x: input
 73 |     :param x_hat: target
 74 |     :param dim_img:
 75 |     :param dim_z:
 76 |     :param n_hidden:
 77 |     :param keep_prob:
 78 |     :return:
 79 |     """
 80 |     batchsz = x.size(0)
 81 |     # encoding
 82 |     mu, log_sigma = encoder(x)
 83 |     sigma = torch.exp(log_sigma)
 84 |     # sampling by re-parameterization technique
 85 |     z = mu + sigma * torch.randn_like(mu)
 86 | 
 87 |     # decoding
 88 |     y = decoder(z)
 89 |     # y = torch.clamp(y, 1e-8, 1 - 1e-8)
 90 | 
 91 | 
 92 |     # loss
 93 |     # marginal_likelihood2 = torch.sum(x_target * torch.log(y) + (1 - x_target) * torch.log(1 - y)) / batchsz
 94 |     # marginal_likelihood = -F.binary_cross_entropy(y, x_target, reduction='sum') / batchsz
 95 |     marginal_likelihood = -torch.pow(x_target - y, 2).sum() / batchsz
 96 |     # print(marginal_likelihood2.item(), marginal_likelihood.item())
 97 | 
 98 |     KL_divergence = 0.5 * torch.sum(
 99 |                                 torch.pow(mu, 2) +
100 |                                 torch.pow(sigma, 2) -
101 |                                 torch.log(1e-8 + torch.pow(sigma, 2)) - 1
102 |                                ).sum() / batchsz
103 | 
104 |     ELBO = marginal_likelihood - KL_divergence
105 | 
106 |     loss = -ELBO
107 | 
108 |     return y, z, loss, marginal_likelihood, KL_divergence


--------------------------------------------------------------------------------