Add a converter which shows how to convert old torch models into pytorch models.

Author: suquark

converter.py

@@ -0,0 +1,130 @@
+import os
+
+import torch
+import torch.nn as nn
+from torch.utils.serialization import load_lua
+
+from models import VGGEncoder, VGGDecoder
+from photo_wct import PhotoWCT
+
+
+def weight_assign(lua, pth, maps):
+    for k, v in maps.items():
+        getattr(pth, k).weight = nn.Parameter(lua.get(v).weight.float())
+        getattr(pth, k).bias = nn.Parameter(lua.get(v).bias.float())
+
+
+def photo_wct_loader(p_wct):
+    p_wct.e1.load_state_dict(torch.load('pth_models/vgg_normalised_conv1.pth'))
+    p_wct.d1.load_state_dict(torch.load('pth_models/feature_invertor_conv1.pth'))
+    p_wct.e2.load_state_dict(torch.load('pth_models/vgg_normalised_conv2.pth'))
+    p_wct.d2.load_state_dict(torch.load('pth_models/feature_invertor_conv2.pth'))
+    p_wct.e3.load_state_dict(torch.load('pth_models/vgg_normalised_conv3.pth'))
+    p_wct.d3.load_state_dict(torch.load('pth_models/feature_invertor_conv3.pth'))
+    p_wct.e4.load_state_dict(torch.load('pth_models/vgg_normalised_conv4.pth'))
+    p_wct.d4.load_state_dict(torch.load('pth_models/feature_invertor_conv4.pth'))
+
+
+if __name__ == '__main__':
+    if not os.path.exists('pth_models'):
+        os.mkdir('pth_models')
+    
+    ## VGGEncoder1
+    vgg1 = load_lua('models/vgg_normalised_conv1_1_mask.t7')
+    e1 = VGGEncoder(1)
+    weight_assign(vgg1, e1, {
+        'conv0': 0,
+        'conv1_1': 2,
+    })
+    torch.save(e1.state_dict(), 'pth_models/vgg_normalised_conv1.pth')
+    
+    ## VGGDecoder1
+    inv1 = load_lua('models/feature_invertor_conv1_1_mask.t7')
+    d1 = VGGDecoder(1)
+    weight_assign(inv1, d1, {
+        'conv1_1': 1,
+    })
+    torch.save(d1.state_dict(), 'pth_models/feature_invertor_conv1.pth')
+    
+    ## VGGEncoder2
+    vgg2 = load_lua('models/vgg_normalised_conv2_1_mask.t7')
+    e2 = VGGEncoder(2)
+    weight_assign(vgg2, e2, {
+        'conv0': 0,
+        'conv1_1': 2,
+        'conv1_2': 5,
+        'conv2_1': 9,
+    })
+    torch.save(e2.state_dict(), 'pth_models/vgg_normalised_conv2.pth')
+    
+    ## VGGDecoder2
+    inv2 = load_lua('models/feature_invertor_conv2_1_mask.t7')
+    d2 = VGGDecoder(2)
+    weight_assign(inv2, d2, {
+        'conv2_1': 1,
+        'conv1_2': 5,
+        'conv1_1': 8,
+    })
+    torch.save(d2.state_dict(), 'pth_models/feature_invertor_conv2.pth')
+    
+    ## VGGEncoder3
+    vgg3 = load_lua('models/vgg_normalised_conv3_1_mask.t7')
+    e3 = VGGEncoder(3)
+    weight_assign(vgg3, e3, {
+        'conv0': 0,
+        'conv1_1': 2,
+        'conv1_2': 5,
+        'conv2_1': 9,
+        'conv2_2': 12,
+        'conv3_1': 16,
+    })
+    torch.save(e3.state_dict(), 'pth_models/vgg_normalised_conv3.pth')
+    
+    ## VGGDecoder3
+    inv3 = load_lua('models/feature_invertor_conv3_1_mask.t7')
+    d3 = VGGDecoder(3)
+    weight_assign(inv3, d3, {
+        'conv3_1': 1,
+        'conv2_2': 5,
+        'conv2_1': 8,
+        'conv1_2': 12,
+        'conv1_1': 15,
+    })
+    torch.save(d3.state_dict(), 'pth_models/feature_invertor_conv3.pth')
+    
+    ## VGGEncoder4
+    vgg4 = load_lua('models/vgg_normalised_conv4_1_mask.t7')
+    e4 = VGGEncoder(4)
+    weight_assign(vgg4, e4, {
+        'conv0': 0,
+        'conv1_1': 2,
+        'conv1_2': 5,
+        'conv2_1': 9,
+        'conv2_2': 12,
+        'conv3_1': 16,
+        'conv3_2': 19,
+        'conv3_3': 22,
+        'conv3_4': 25,
+        'conv4_1': 29,
+    })
+    torch.save(e4.state_dict(), 'pth_models/vgg_normalised_conv4.pth')
+    
+    ## VGGDecoder4
+    inv4 = load_lua('models/feature_invertor_conv4_1_mask.t7')
+    d4 = VGGDecoder(4)
+    weight_assign(inv4, d4, {
+        'conv4_1': 1,
+        'conv3_4': 5,
+        'conv3_3': 8,
+        'conv3_2': 11,
+        'conv3_1': 14,
+        'conv2_2': 18,
+        'conv2_1': 21,
+        'conv1_2': 25,
+        'conv1_1': 28,
+    })
+    torch.save(d4.state_dict(), 'pth_models/feature_invertor_conv4.pth')
+    
+    p_wct = PhotoWCT()
+    photo_wct_loader(p_wct)
+    torch.save(p_wct.state_dict(), 'PhotoWCTModels/photo_wct.pth')