├── .DS_Store
├── .idea
    ├── WAE.iml
    ├── deployment.xml
    ├── dictionaries
    │   └── chensnathan.xml
    ├── misc.xml
    ├── modules.xml
    ├── vcs.xml
    ├── webServers.xml
    └── workspace.xml
├── README.md
├── images
    ├── .DS_Store
    ├── real_image_100.png
    ├── recon_image_100.png
    └── sample_image_100.png
└── wae_for_mnist.py


/.DS_Store:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsnedy/WAE_Pytorch/c599aef1cbc6a73210b2fe045998b55bced22f01/.DS_Store


--------------------------------------------------------------------------------
/.idea/WAE.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="py.test" />
10 |     <option name="PROJECT_TEST_RUNNER" value="py.test" />
11 |   </component>
12 | </module>


--------------------------------------------------------------------------------
/.idea/deployment.xml:
--------------------------------------------------------------------------------
 1 | <?xml version="1.0" encoding="UTF-8"?>
 2 | <project version="4">
 3 |   <component name="PublishConfigData" serverName="wae">
 4 |     <serverData>
 5 |       <paths name="wae">
 6 |         <serverdata>
 7 |           <mappings>
 8 |             <mapping deploy="/WAE" local="$PROJECT_DIR$" web="/" />
 9 |           </mappings>
10 |         </serverdata>
11 |       </paths>
12 |     </serverData>
13 |   </component>
14 | </project>


--------------------------------------------------------------------------------
/.idea/dictionaries/chensnathan.xml:
--------------------------------------------------------------------------------
1 | <component name="ProjectDictionaryState">
2 |   <dictionary name="chensnathan" />
3 | </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 2.7 (PycharmProjects)" 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/WAE.iml" filepath="$PROJECT_DIR$/.idea/WAE.iml" />
6 |     </modules>
7 |   </component>
8 | </project>


--------------------------------------------------------------------------------
/.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/webServers.xml:
--------------------------------------------------------------------------------
 1 | <?xml version="1.0" encoding="UTF-8"?>
 2 | <project version="4">
 3 |   <component name="WebServers">
 4 |     <option name="servers">
 5 |       <webServer id="b3f0b8ec-552d-4ed3-954b-ff447799edea" name="wae" url="http://159.226.176.178">
 6 |         <fileTransfer host="159.226.176.178" port="22" rootFolder="/home/qchen" accessType="SFTP">
 7 |           <advancedOptions>
 8 |             <advancedOptions dataProtectionLevel="Private" />
 9 |           </advancedOptions>
10 |           <option name="port" value="22" />
11 |         </fileTransfer>
12 |       </webServer>
13 |     </option>
14 |   </component>
15 | </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="aca82054-ac06-4107-8d58-f51bf2b102b1" name="Default" comment="">
  5 |       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" afterPath="$PROJECT_DIR$/.idea/workspace.xml" />
  6 |       <change beforePath="$PROJECT_DIR$/README.md" afterPath="$PROJECT_DIR$/README.md" />
  7 |       <change beforePath="$PROJECT_DIR$/wae_for_mnist.py" afterPath="$PROJECT_DIR$/wae_for_mnist.py" />
  8 |     </list>
  9 |     <option name="EXCLUDED_CONVERTED_TO_IGNORED" value="true" />
 10 |     <option name="TRACKING_ENABLED" value="true" />
 11 |     <option name="SHOW_DIALOG" value="false" />
 12 |     <option name="HIGHLIGHT_CONFLICTS" value="true" />
 13 |     <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
 14 |     <option name="LAST_RESOLUTION" value="IGNORE" />
 15 |   </component>
 16 |   <component name="FileEditorManager">
 17 |     <leaf SIDE_TABS_SIZE_LIMIT_KEY="300">
 18 |       <file leaf-file-name="wae_for_mnist.py" pinned="false" current-in-tab="false">
 19 |         <entry file="file://$PROJECT_DIR$/wae_for_mnist.py">
 20 |           <provider selected="true" editor-type-id="text-editor">
 21 |             <state relative-caret-position="315">
 22 |               <caret line="15" column="38" lean-forward="true" selection-start-line="15" selection-start-column="38" selection-end-line="15" selection-end-column="38" />
 23 |               <folding>
 24 |                 <element signature="e#0#12#0" expanded="true" />
 25 |               </folding>
 26 |             </state>
 27 |           </provider>
 28 |         </entry>
 29 |       </file>
 30 |       <file leaf-file-name="README.md" pinned="false" current-in-tab="true">
 31 |         <entry file="file://$PROJECT_DIR$/README.md">
 32 |           <provider selected="true" editor-type-id="split-provider[text-editor;markdown-preview-editor]">
 33 |             <state split_layout="SPLIT">
 34 |               <first_editor relative-caret-position="147">
 35 |                 <caret line="7" column="31" lean-forward="false" selection-start-line="7" selection-start-column="31" selection-end-line="7" selection-end-column="31" />
 36 |                 <folding />
 37 |               </first_editor>
 38 |               <second_editor />
 39 |             </state>
 40 |           </provider>
 41 |         </entry>
 42 |       </file>
 43 |     </leaf>
 44 |   </component>
 45 |   <component name="FileTemplateManagerImpl">
 46 |     <option name="RECENT_TEMPLATES">
 47 |       <list>
 48 |         <option value="Python Script" />
 49 |       </list>
 50 |     </option>
 51 |   </component>
 52 |   <component name="FindInProjectRecents">
 53 |     <findStrings>
 54 |       <find>normal</find>
 55 |     </findStrings>
 56 |   </component>
 57 |   <component name="Git.Settings">
 58 |     <option name="RECENT_GIT_ROOT_PATH" value="$PROJECT_DIR$" />
 59 |   </component>
 60 |   <component name="IdeDocumentHistory">
 61 |     <option name="CHANGED_PATHS">
 62 |       <list>
 63 |         <option value="$PROJECT_DIR$/wae_for_mnist.py" />
 64 |         <option value="$PROJECT_DIR$/README.md" />
 65 |       </list>
 66 |     </option>
 67 |   </component>
 68 |   <component name="JsBuildToolGruntFileManager" detection-done="true" sorting="DEFINITION_ORDER" />
 69 |   <component name="JsBuildToolPackageJson" detection-done="true" sorting="DEFINITION_ORDER" />
 70 |   <component name="JsGulpfileManager">
 71 |     <detection-done>true</detection-done>
 72 |     <sorting>DEFINITION_ORDER</sorting>
 73 |   </component>
 74 |   <component name="ProjectFrameBounds" fullScreen="true">
 75 |     <option name="width" value="1440" />
 76 |     <option name="height" value="900" />
 77 |   </component>
 78 |   <component name="ProjectInspectionProfilesVisibleTreeState">
 79 |     <entry key="Project Default">
 80 |       <profile-state>
 81 |         <expanded-state>
 82 |           <State>
 83 |             <id />
 84 |           </State>
 85 |           <State>
 86 |             <id>CSS</id>
 87 |           </State>
 88 |           <State>
 89 |             <id>Probable bugsCSS</id>
 90 |           </State>
 91 |           <State>
 92 |             <id>RELAX NG</id>
 93 |           </State>
 94 |         </expanded-state>
 95 |         <selected-state>
 96 |           <State>
 97 |             <id>AngularJS</id>
 98 |           </State>
 99 |         </selected-state>
100 |       </profile-state>
101 |     </entry>
102 |   </component>
103 |   <component name="ProjectLevelVcsManager" settingsEditedManually="true" />
104 |   <component name="ProjectView">
105 |     <navigator currentView="ProjectPane" proportions="" version="1">
106 |       <flattenPackages />
107 |       <showMembers />
108 |       <showModules />
109 |       <showLibraryContents />
110 |       <hideEmptyPackages />
111 |       <abbreviatePackageNames />
112 |       <autoscrollToSource />
113 |       <autoscrollFromSource />
114 |       <sortByType />
115 |       <manualOrder />
116 |       <foldersAlwaysOnTop value="true" />
117 |     </navigator>
118 |     <panes>
119 |       <pane id="Scratches" />
120 |       <pane id="Scope" />
121 |       <pane id="ProjectPane">
122 |         <subPane>
123 |           <expand>
124 |             <path>
125 |               <item name="WAE" type="b2602c69:ProjectViewProjectNode" />
126 |               <item name="WAE" type="462c0819:PsiDirectoryNode" />
127 |             </path>
128 |           </expand>
129 |           <select />
130 |         </subPane>
131 |       </pane>
132 |     </panes>
133 |   </component>
134 |   <component name="PropertiesComponent">
135 |     <property name="settings.editor.selected.configurable" value="com.jetbrains.python.configuration.PyActiveSdkModuleConfigurable" />
136 |     <property name="WebServerToolWindowFactoryState" value="true" />
137 |     <property name="WebServerToolWindowPanel.toolwindow.show.size" value="false" />
138 |     <property name="WebServerToolWindowPanel.toolwindow.show.date" value="false" />
139 |     <property name="WebServerToolWindowPanel.toolwindow.show.permissions" value="false" />
140 |     <property name="WebServerToolWindowPanel.toolwindow.highlight.symlinks" value="true" />
141 |     <property name="WebServerToolWindowPanel.toolwindow.highlight.mappings" value="true" />
142 |   </component>
143 |   <component name="RunDashboard">
144 |     <option name="ruleStates">
145 |       <list>
146 |         <RuleState>
147 |           <option name="name" value="ConfigurationTypeDashboardGroupingRule" />
148 |         </RuleState>
149 |         <RuleState>
150 |           <option name="name" value="StatusDashboardGroupingRule" />
151 |         </RuleState>
152 |       </list>
153 |     </option>
154 |   </component>
155 |   <component name="ShelveChangesManager" show_recycled="false">
156 |     <option name="remove_strategy" value="false" />
157 |   </component>
158 |   <component name="SvnConfiguration">
159 |     <configuration />
160 |   </component>
161 |   <component name="TaskManager">
162 |     <task active="true" id="Default" summary="Default task">
163 |       <changelist id="aca82054-ac06-4107-8d58-f51bf2b102b1" name="Default" comment="" />
164 |       <created>1517468803468</created>
165 |       <option name="number" value="Default" />
166 |       <option name="presentableId" value="Default" />
167 |       <updated>1517468803468</updated>
168 |     </task>
169 |     <servers />
170 |   </component>
171 |   <component name="ToolWindowManager">
172 |     <frame x="0" y="0" width="1440" height="900" extended-state="0" />
173 |     <editor active="true" />
174 |     <layout>
175 |       <window_info id="TODO" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="6" side_tool="false" content_ui="tabs" />
176 |       <window_info id="Event Log" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="7" side_tool="true" content_ui="tabs" />
177 |       <window_info id="Version Control" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="7" side_tool="false" content_ui="tabs" />
178 |       <window_info id="Python Console" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="7" side_tool="false" content_ui="tabs" />
179 |       <window_info id="Run" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="2" side_tool="false" content_ui="tabs" />
180 |       <window_info id="Terminal" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="7" side_tool="false" content_ui="tabs" />
181 |       <window_info id="Remote Host" active="false" anchor="right" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.3298611" sideWeight="0.5" order="3" side_tool="false" content_ui="tabs" />
182 |       <window_info id="Project" active="false" anchor="left" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="true" show_stripe_button="true" weight="0.16597222" sideWeight="0.5" order="0" side_tool="false" content_ui="combo" />
183 |       <window_info id="Docker" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="false" weight="0.33" sideWeight="0.5" order="7" side_tool="false" content_ui="tabs" />
184 |       <window_info id="Database" active="false" anchor="right" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="3" side_tool="false" content_ui="tabs" />
185 |       <window_info id="SciView" active="false" anchor="right" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="3" side_tool="false" content_ui="tabs" />
186 |       <window_info id="Structure" active="false" anchor="left" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.25" sideWeight="0.5" order="1" side_tool="false" content_ui="tabs" />
187 |       <window_info id="Favorites" active="false" anchor="left" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="2" side_tool="true" content_ui="tabs" />
188 |       <window_info id="Debug" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.4" sideWeight="0.5" order="3" side_tool="false" content_ui="tabs" />
189 |       <window_info id="Cvs" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.25" sideWeight="0.5" order="4" side_tool="false" content_ui="tabs" />
190 |       <window_info id="Message" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="0" side_tool="false" content_ui="tabs" />
191 |       <window_info id="Commander" active="false" anchor="right" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.4" sideWeight="0.5" order="0" side_tool="false" content_ui="tabs" />
192 |       <window_info id="Inspection" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.4" sideWeight="0.5" order="5" side_tool="false" content_ui="tabs" />
193 |       <window_info id="File Transfer" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="true" show_stripe_button="true" weight="0.32941177" sideWeight="0.5" order="7" side_tool="false" content_ui="tabs" />
194 |       <window_info id="Hierarchy" active="false" anchor="right" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.25" sideWeight="0.5" order="2" side_tool="false" content_ui="combo" />
195 |       <window_info id="Find" active="false" anchor="bottom" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.33" sideWeight="0.5" order="1" side_tool="false" content_ui="tabs" />
196 |       <window_info id="Ant Build" active="false" anchor="right" auto_hide="false" internal_type="DOCKED" type="DOCKED" visible="false" show_stripe_button="true" weight="0.25" sideWeight="0.5" order="1" side_tool="false" content_ui="tabs" />
197 |     </layout>
198 |   </component>
199 |   <component name="TypeScriptGeneratedFilesManager">
200 |     <option name="version" value="1" />
201 |   </component>
202 |   <component name="VcsContentAnnotationSettings">
203 |     <option name="myLimit" value="2678400000" />
204 |   </component>
205 |   <component name="XDebuggerManager">
206 |     <breakpoint-manager />
207 |     <watches-manager />
208 |   </component>
209 |   <component name="editorHistoryManager">
210 |     <entry file="file://$PROJECT_DIR$/wae_for_mnist.py">
211 |       <provider selected="true" editor-type-id="text-editor">
212 |         <state relative-caret-position="6216">
213 |           <caret line="296" column="0" lean-forward="true" selection-start-line="296" selection-start-column="0" selection-end-line="296" selection-end-column="0" />
214 |           <folding>
215 |             <element signature="e#0#12#0" expanded="true" />
216 |           </folding>
217 |         </state>
218 |       </provider>
219 |     </entry>
220 |     <entry file="file://$PROJECT_DIR$/README.md">
221 |       <provider selected="true" editor-type-id="split-provider[text-editor;markdown-preview-editor]">
222 |         <state split_layout="SPLIT">
223 |           <first_editor relative-caret-position="0">
224 |             <caret line="0" column="0" lean-forward="false" selection-start-line="0" selection-start-column="0" selection-end-line="0" selection-end-column="0" />
225 |             <folding />
226 |           </first_editor>
227 |           <second_editor />
228 |         </state>
229 |       </provider>
230 |     </entry>
231 |     <entry file="file://$PROJECT_DIR$/wae_for_mnist.py">
232 |       <provider selected="true" editor-type-id="text-editor">
233 |         <state relative-caret-position="0">
234 |           <caret line="0" column="0" lean-forward="false" selection-start-line="0" selection-start-column="0" selection-end-line="0" selection-end-column="0" />
235 |           <folding>
236 |             <element signature="e#0#12#0" expanded="true" />
237 |           </folding>
238 |         </state>
239 |       </provider>
240 |     </entry>
241 |     <entry file="file://$PROJECT_DIR$/wae_for_mnist.py">
242 |       <provider selected="true" editor-type-id="text-editor">
243 |         <state relative-caret-position="315">
244 |           <caret line="15" column="38" lean-forward="true" selection-start-line="15" selection-start-column="38" selection-end-line="15" selection-end-column="38" />
245 |           <folding>
246 |             <element signature="e#0#12#0" expanded="true" />
247 |           </folding>
248 |         </state>
249 |       </provider>
250 |     </entry>
251 |     <entry file="file://$PROJECT_DIR$/README.md">
252 |       <provider selected="true" editor-type-id="split-provider[text-editor;markdown-preview-editor]">
253 |         <state split_layout="SPLIT">
254 |           <first_editor relative-caret-position="147">
255 |             <caret line="7" column="31" lean-forward="false" selection-start-line="7" selection-start-column="31" selection-end-line="7" selection-end-column="31" />
256 |             <folding />
257 |           </first_editor>
258 |           <second_editor />
259 |         </state>
260 |       </provider>
261 |       <provider editor-type-id="text-editor">
262 |         <state relative-caret-position="147">
263 |           <caret line="7" column="0" lean-forward="true" selection-start-line="7" selection-start-column="0" selection-end-line="7" selection-end-column="0" />
264 |           <folding />
265 |         </state>
266 |       </provider>
267 |     </entry>
268 |   </component>
269 |   <component name="masterDetails">
270 |     <states>
271 |       <state key="ScopeChooserConfigurable.UI">
272 |         <settings>
273 |           <splitter-proportions>
274 |             <option name="proportions">
275 |               <list>
276 |                 <option value="0.2" />
277 |               </list>
278 |             </option>
279 |           </splitter-proportions>
280 |         </settings>
281 |       </state>
282 |     </states>
283 |   </component>
284 | </project>


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Pytorch implementation of Wasserstein Auto-Encoders
 2 | The reference of this code is the original implementation in TensorFlow:
 3 | <https://github.com/tolstikhin/wae/>
 4 | 
 5 | 
 6 | ##### The results of epoch 100 are as follow:
 7 | ##### real images
 8 | ![real_image](images/real_image_100.png) 
 9 | ##### recon images
10 | ![recon_image](images/recon_image_100.png) 
11 | ##### sample images
12 | ![sample_image](images/sample_image_100.png)
13 | 


--------------------------------------------------------------------------------
/images/.DS_Store:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsnedy/WAE_Pytorch/c599aef1cbc6a73210b2fe045998b55bced22f01/images/.DS_Store


--------------------------------------------------------------------------------
/images/real_image_100.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsnedy/WAE_Pytorch/c599aef1cbc6a73210b2fe045998b55bced22f01/images/real_image_100.png


--------------------------------------------------------------------------------
/images/recon_image_100.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsnedy/WAE_Pytorch/c599aef1cbc6a73210b2fe045998b55bced22f01/images/recon_image_100.png


--------------------------------------------------------------------------------
/images/sample_image_100.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsnedy/WAE_Pytorch/c599aef1cbc6a73210b2fe045998b55bced22f01/images/sample_image_100.png


--------------------------------------------------------------------------------
/wae_for_mnist.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn as nn
  3 | import torch.nn.functional as F
  4 | from torch.autograd import Variable
  5 | import torch.optim as optim
  6 | from torchvision.datasets import MNIST
  7 | from torch.utils.data import DataLoader
  8 | import torchvision
  9 | import itertools
 10 | from torchvision.utils import save_image
 11 | 
 12 | 
 13 | class Encoder(nn.Module):
 14 |     def __init__(self, in_channels, num_filters, num_layers, z_size, is_training=False):
 15 |         super(Encoder, self).__init__()
 16 |         self.in_channels = in_channels
 17 |         self.is_training = is_training
 18 |         conv_list = []
 19 |         for i in xrange(num_layers):
 20 |             scale = 2 ** (num_layers - i - 1)
 21 |             conv_i = nn.Sequential(
 22 |                 nn.Conv2d(
 23 |                     in_channels=self.in_channels, out_channels=num_filters / scale,
 24 |                     kernel_size=4, stride=2, padding=2
 25 |                 ),
 26 |                 nn.BatchNorm2d(num_features=num_filters / scale),
 27 |                 nn.ReLU(inplace=True)
 28 |             )
 29 |             conv_list.append(conv_i)
 30 |             self.in_channels = num_filters / scale
 31 |         self.conv = nn.Sequential(*conv_list)
 32 |         self.linear = nn.Linear(in_features=3 * 3 * num_filters, out_features=z_size)
 33 |         # initialize weights
 34 |         for m in self.modules():
 35 |             if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
 36 |                 m.weight.data.normal_(0.0, 0.0099999)
 37 |                 m.bias.data.zero_()
 38 |             elif isinstance(m, nn.BatchNorm2d):
 39 |                 m.weight.data.normal_(0.0, 0.01)
 40 |                 m.bias.data.zero_()
 41 | 
 42 |     def forward(self, x):
 43 |         if self.is_training:
 44 |             noise = torch.normal(torch.zeros_like(x.data), std=0.01)
 45 |             x.data += noise
 46 |         conv_out = self.conv(x).view(-1, 3 * 3 * 1024)
 47 |         z = self.linear(conv_out)
 48 |         return z
 49 | 
 50 | 
 51 | class Decoder(nn.Module):
 52 |     def __init__(self, num_filters, num_layers, z_size, output_shape):
 53 |         super(Decoder, self).__init__()
 54 |         height = output_shape / 2 ** (num_layers - 1) + 1
 55 |         width = output_shape / 2 ** (num_layers - 1) + 1
 56 |         self.linear1 = nn.Sequential(
 57 |             nn.Linear(in_features=z_size, out_features=num_filters * height * width),
 58 |             nn.ReLU(inplace=True)
 59 |         )
 60 |         self.in_channels = num_filters
 61 |         deconv_list = []
 62 |         for i in xrange(num_layers - 1):
 63 |             scale = 2 ** (i + 1)
 64 |             deconv_i = nn.Sequential(
 65 |                 nn.ConvTranspose2d(
 66 |                     in_channels=self.in_channels, out_channels=num_filters / scale,
 67 |                     kernel_size=4, stride=2, padding=2, output_padding=1
 68 |                 ),
 69 |                 nn.BatchNorm2d(num_features=num_filters / scale),
 70 |                 nn.ReLU(inplace=True)
 71 |             )
 72 |             deconv_list.append(deconv_i)
 73 |             self.in_channels = num_filters / scale
 74 |         self.deconv = nn.Sequential(*deconv_list)
 75 |         self.deconv_last = nn.ConvTranspose2d(
 76 |             in_channels=self.in_channels, out_channels=1,
 77 |             kernel_size=4, stride=1, padding=2
 78 |         )
 79 |         # initialize weights
 80 |         for m in self.modules():
 81 |             if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):
 82 |                 m.weight.data.normal_(0.0, 0.0099999)
 83 |                 m.bias.data.zero_()
 84 |             elif isinstance(m, nn.BatchNorm2d):
 85 |                 m.weight.data.normal_(0.0, 0.01)
 86 |                 m.bias.data.zero_()
 87 | 
 88 |     def forward(self, z):
 89 |         linear_out = self.linear1(z)
 90 |         deconv_input = linear_out.view(-1, 1024, 8, 8)
 91 |         deconv_out = self.deconv(deconv_input)
 92 |         recon_x = self.deconv_last(deconv_out)
 93 |         return F.sigmoid(recon_x), recon_x
 94 | 
 95 | 
 96 | class Adversary_z(nn.Module):
 97 |     def __init__(self, num_filters, num_layers, z_size):
 98 |         super(Adversary_z, self).__init__()
 99 |         self.in_features = z_size
100 |         linears = []
101 |         for i in xrange(num_layers):
102 |             linear_i = nn.Sequential(
103 |                 nn.Linear(in_features=self.in_features, out_features=num_filters),
104 |                 nn.ReLU(inplace=True)
105 |             )
106 |             linears.append(linear_i)
107 |             self.in_features = num_filters
108 |         self.linear = nn.Sequential(*linears)
109 |         self.final_linear = nn.Linear(in_features=self.in_features, out_features=1)
110 |         # initialize weights
111 |         for m in self.modules():
112 |             if isinstance(m, nn.Linear):
113 |                 m.weight.data.normal_(0.0, 0.0099999)
114 |                 m.bias.data.zero_()
115 | 
116 |     def forward(self, z):
117 |         linear_out = self.linear(z)
118 |         out = self.final_linear(linear_out)
119 |         return out
120 | 
121 | 
122 | def run():
123 |     # define the hyper-parameters
124 |     batch_size = 100
125 |     e_pretrain_batch_size = 1000
126 |     pretrain_epochs = 200
127 |     epochs = 100
128 |     z_size = 8
129 |     lam = 10
130 | 
131 |     e_num_filters = 1024
132 |     e_num_layers = 4
133 | 
134 |     g_num_filters = 1024
135 |     g_num_layers = 3
136 | 
137 |     d_num_filters = 512
138 |     d_num_layers = 4
139 | 
140 |     # download and load data
141 |     train_data = MNIST(
142 |         root='data/',
143 |         train=True,
144 |         transform=torchvision.transforms.ToTensor()
145 |     )
146 |     test_data = MNIST(
147 |         root='data/',
148 |         train=False,
149 |         transform=torchvision.transforms.ToTensor()
150 |     )
151 |     # get the data batch
152 |     pretrain_data_loader = DataLoader(dataset=train_data, batch_size=e_pretrain_batch_size, shuffle=True)
153 |     pretest_data_loader = DataLoader(dataset=test_data, batch_size=e_pretrain_batch_size, shuffle=False)
154 |     train_data_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)
155 |     test_data_loader = DataLoader(dataset=test_data, batch_size=batch_size, shuffle=False)
156 | 
157 |     # models
158 |     encoder = Encoder(
159 |         in_channels=1, num_filters=e_num_filters, num_layers=e_num_layers, z_size=z_size, is_training=True
160 |     )
161 |     decoder = Decoder(
162 |         num_filters=g_num_filters, num_layers=g_num_layers, z_size=z_size, output_shape=28
163 |     )
164 |     discriminator = Adversary_z(num_filters=d_num_filters, num_layers=d_num_layers, z_size=z_size)
165 | 
166 |     # load model parameters
167 |     # encoder.load_state_dict(torch.load('Pretrain_Encoder_epoch200.pth'))
168 |     # encoder.load_state_dict(torch.load('encoder_100.pth'))
169 |     # decoder.load_state_dict(torch.load('decoder_100.pth'))
170 |     # discriminator.load_state_dict(torch.load('discriminator_100.pth'))
171 | 
172 |     if torch.cuda.is_available():
173 |         encoder, decoder, discriminator = encoder.cuda(), decoder.cuda(), discriminator.cuda()
174 | 
175 |     # define the optimizer
176 |     e_params = encoder.parameters()
177 |     ae_params = itertools.chain(encoder.parameters(), decoder.parameters())
178 |     d_params = discriminator.parameters()
179 |     optimizer_e = optim.Adam(e_params, lr=1e-03, betas=(0.5, 0.999))
180 |     optimizer_ae = optim.Adam(ae_params, lr=1e-03, betas=(0.5, 0.999))
181 |     optimizer_d = optim.Adam(d_params, lr=5e-04, betas=(0.5, 0.999))
182 | 
183 |     # pretrain model
184 |     def sample_pz(batch_size=100, z_size=8):
185 |         return Variable(torch.normal(torch.zeros(batch_size, z_size), std=1).cuda())
186 | 
187 |     def pretrain_loss(encoded, sample_noise):
188 |         # for mean
189 |         mean_qz = torch.mean(encoded, dim=0, keepdim=True)
190 |         mean_pz = torch.mean(sample_noise, dim=0, keepdim=True)
191 |         mean_loss = F.mse_loss(mean_qz, mean_pz)
192 | 
193 |         # for covariance
194 |         cov_qz = torch.matmul((encoded - mean_qz).transpose(0, 1), encoded - mean_qz)
195 |         cov_qz /= e_pretrain_batch_size - 1.
196 |         cov_pz = torch.matmul((sample_noise - mean_pz).transpose(0, 1), sample_noise - mean_pz)
197 |         cov_pz /= e_pretrain_batch_size - 1.
198 |         cov_loss = F.mse_loss(cov_qz, cov_pz)
199 | 
200 |         return mean_loss + cov_loss
201 | 
202 |     def encoder_pretrain(epoch):
203 |         encoder.train()
204 |         for batch_idx, (data, _) in enumerate(pretrain_data_loader):
205 |             if torch.cuda.is_available():
206 |                 data = data.cuda()
207 |             data = Variable(data)
208 |             sample_noise = sample_pz(e_pretrain_batch_size, z_size)
209 |             encoded = encoder(data)
210 | 
211 |             optimizer_e.zero_grad()
212 |             loss_pretrain = pretrain_loss(encoded, sample_noise)
213 |             loss_pretrain.backward()
214 |             optimizer_e.step()
215 |             if batch_idx % 10 == 0:
216 |                 print("Train Epoch: {} [{}/{} ({:.0f}%)] \tLoss: {:.6f}".format(
217 |                     epoch, batch_idx * len(data), len(pretrain_data_loader.dataset),
218 |                            100. * batch_idx / len(pretrain_data_loader), loss_pretrain.data[0]
219 |                 ))
220 |         # save the pretrain model at the last epoch
221 |         if epoch == 200:
222 |             torch.save(encoder.state_dict(), 'Pretrain_Encoder_epoch{:02d}.pth'.format(epoch))
223 | 
224 |     # train models
225 |     d_loss_function = nn.BCEWithLogitsLoss()
226 | 
227 |     def gan_loss(sample_qz, sample_pz):
228 |         logits_qz = discriminator(sample_qz)
229 |         logits_pz = discriminator(sample_pz)
230 | 
231 |         # losses
232 |         loss_qz = d_loss_function(logits_qz, torch.zeros_like(logits_qz))
233 |         loss_pz = d_loss_function(logits_pz, torch.ones_like(logits_pz))
234 |         loss_qz_trick = d_loss_function(logits_qz, torch.ones_like(logits_qz))
235 |         loss_adversary = lam * (loss_qz + loss_pz)
236 |         loss_penalty = loss_qz_trick
237 |         return (loss_adversary, logits_qz, logits_pz), loss_penalty
238 | 
239 |     def train(epoch):
240 |         for param_group in optimizer_ae.param_groups:
241 |             print(param_group['lr'], "learning rate for Auto-Encoder.")
242 |         for param_group in optimizer_d.param_groups:
243 |             print(param_group['lr'], "learning rate for Discriminator.")
244 |         encoder.train(), decoder.train(), discriminator.train()
245 |         for batch_idx, (data, _) in enumerate(train_data_loader):
246 |             if torch.cuda.is_available():
247 |                 data = data.cuda()
248 |             data = Variable(data)
249 |             sample_noise = sample_pz(batch_size, z_size)
250 | 
251 |             encoded = encoder(data)
252 |             # for reconstructed
253 |             recon_x, recon_logits = decoder(encoded)
254 |             # for sample
255 |             decoded, decoded_logits = decoder(sample_noise)
256 | 
257 |             # losses
258 |             recon_loss = F.mse_loss(recon_x, data)
259 |             loss_gan, loss_penalty = gan_loss(encoded, sample_noise)
260 |             loss_wae = recon_loss + lam * loss_penalty
261 |             loss_adv = loss_gan[0]
262 | 
263 |             # optimize wae
264 |             encoder.zero_grad()
265 |             decoder.zero_grad()
266 |             loss_wae.backward(retain_graph=True)
267 |             optimizer_ae.step()
268 | 
269 |             # optimize adv
270 |             discriminator.zero_grad()
271 |             loss_adv.backward()
272 |             optimizer_d.step()
273 | 
274 |             if batch_idx % 10 == 0:
275 |                 print("Train Epoch: {} [{}/{} ({:.0f}%)] \tWAE_Loss: {:.6f}\tD_Loss: {:.6f}".format(
276 |                     epoch, batch_idx * len(data), len(train_data_loader.dataset),
277 |                            100. * batch_idx / len(train_data_loader), loss_wae.data[0], loss_adv.data[0]
278 |                 ))
279 |         # save images and save models
280 |         save_image(data.cpu().data, 'real_image_{:02d}.png'.format(epoch), nrow=10)
281 |         save_image(recon_x.cpu().data, 'recon_image_{:02d}.png'.format(epoch), nrow=10)
282 |         save_image(decoded.cpu().data, 'sample_image_{:02d}.png'.format(epoch), nrow=10)
283 |         if epoch % 50 == 0:
284 |             torch.save(encoder.state_dict(), 'encoder_{:02d}.pth'.format(epoch))
285 |             torch.save(decoder.state_dict(), 'decoder_{:02d}.pth'.format(epoch))
286 |             torch.save(discriminator.state_dict(), 'discriminator_{:02d}.pth'.format(epoch))
287 | 
288 |     # add the learning rate adjust function
289 |     def adjust_learning_rate_manual(optimizer, epoch):
290 |         for param_group in optimizer.param_groups:
291 |             if epoch == 30:
292 |                 param_group['lr'] /= 2.
293 |             elif epoch == 50:
294 |                 param_group['lr'] /= 2.5
295 |             elif epoch == 100:
296 |                 param_group['lr'] /= 2.
297 | 
298 |     print("=========> Pretrain encoder")
299 |     for epoch in range(1, pretrain_epochs+1):
300 |         encoder_pretrain(epoch)
301 |     print("=========> Train models")
302 |     for epoch in range(1, epochs+1):
303 |         # do not need to use adjust_learning_rate in mnist
304 |         # adjust_learning_rate_manual(optimizer_ae, epoch)
305 |         # adjust_learning_rate_manual(optimizer_d, epoch)
306 |         train(epoch)
307 | 
308 | 
309 | if __name__ == "__main__":
310 |     run()
311 | 


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