关于TensorFlow代码的PyTorch等效性的问题

Question

我正在尝试使用PyTorch框架重新实现TensorFlow网络。

问题

给定形状为（9x15）的特征向量z，我想训练生成器G，该生成器输出两个sett值。

1. 一个（9x5）矩阵，f。
2. 一个（4x9x9）矩阵，W。

TensorFlow实现

以下是可用的TensorFlow代码。

dim_atom = 9
dim_bond_type = 4
dim_atom_type = 5
dim_final_feature = 15
dim_final_1 = dim_atom_type
dim_final_2 = dim_atom * dim_bond_type * 15

X = tf.placeholder(tf.float32, shape=[None, dim_atom, dim_final_feature])
W1 = tf.Variable(xavier_init([dim_atom * dim_final_feature, 128]))
b1 = tf.Variable(tf.zeros(shape=[128]))
W11 = tf.Variable(xavier_init([128, 256]))
b11 = tf.Variable(tf.zeros(shape=[256]))
W12 = tf.Variable(xavier_init([256, 512]))
b12 = tf.Variable(tf.zeros(shape=[512]))
W13 = tf.Variable(xavier_init([512, dim_atom * dim_final_1]))
b13 = tf.Variable(tf.zeros(shape=[dim_atom * dim_final_1]))

W2 = tf.Variable(xavier_init([dim_atom * dim_final_feature, 128]))
b2 = tf.Variable(tf.zeros(shape=[128]))
W21 = tf.Variable(xavier_init([128, 256]))
b21 = tf.Variable(tf.zeros(shape=[256]))
W22 = tf.Variable(xavier_init([256, 512]))
b22 = tf.Variable(tf.zeros(shape=[512]))
W23 = tf.Variable(xavier_init([512, dim_final_2]))
b23 = tf.Variable(tf.zeros(shape=[dim_final_2]))

theta = [W1, b1, W11, b11, W12, b12, W13, b13, 
         W2, b2, W21, b21, W22, b22, W23, b23]

def fcn(x):
    out1 = tf.reshape(x, (-1, dim_atom * dim_final_feature))
    out1 = leaky_relu( tf.matmul(out1, W1) + b1 )
    out1 = leaky_relu( tf.matmul(out1, W11) + b11 )
    out1 = leaky_relu( tf.matmul(out1, W12) + b12 )
    out1 = leaky_relu( tf.matmul(out1, W13) + b13 )
    out1 = tf.reshape(out1, (-1, dim_atom, dim_final_1))


    out2 = tf.reshape(x, (-1, dim_atom * dim_final_feature))
    out2 = leaky_relu( tf.matmul(out2, W2) + b2 )
    out2 = leaky_relu( tf.matmul(out2, W21) + b21 )
    out2 = leaky_relu( tf.matmul(out2, W22) + b22 )
    out2 = leaky_relu( tf.matmul(out2, W23) + b23 )
    out2 = tf.reshape(out2, [-1, dim_atom, dim_bond_type, 15])
    out2 = leaky_relu( tf.matmul(tf.transpose(out2, perm=[0,2,1,3]), tf.transpose(out2, perm=[0,2,3,1])) )
    out2 = tf.transpose(out2, perm=[0,2,3,1])
    return [out1, out2]


Y_adj = tf.placeholder(tf.float32, shape=[None, dim_atom, dim_atom, dim_bond_type])
Y_features = tf.placeholder(tf.float32, shape=[None, dim_atom, dim_atom_type])


fcn_loss_1 = tf.nn.softmax_cross_entropy_with_logits(labels=Y_features, logits=fcn(X)[0])
fcn_loss_2 = tf.nn.softmax_cross_entropy_with_logits(labels=Y_adj, logits=fcn(X)[1])
fcn_loss_2 = tf.matrix_band_part(fcn_loss_2,0,-1) - tf.matrix_band_part(fcn_loss_2,0,0)

fcn_loss = tf.reduce_mean(fcn_loss_1) + 2 * tf.reduce_mean(fcn_loss_2)

fcn_solver = (tf.train.AdamOptimizer(learning_rate=0.001)
            .minimize(fcn_loss, var_list=theta))


train_adj_array = to_categorical(np.asarray(train_adj), num_classes=dim_bond_type)
train_features_array = np.transpose(np.asarray(train_features), axes=[0,2,1])


random_idx = list(range(num_train))
shuffle(random_idx)

feature_final = feature_final[random_idx]
train_adj_array = train_adj_array[random_idx]
train_features_array = train_features_array[random_idx]

sess = tf.Session() 
sess.run(tf.global_variables_initializer())

num_epoch = 300
for it in range(num_epoch):    
    for i_batch in range(round(num_train/num_epoch)+1):
        train_sample = feature_final[i_batch * num_epoch : (i_batch+1) * num_epoch] 
        train_adj_sample = train_adj_array[i_batch * num_epoch : (i_batch+1) * num_epoch] 
        train_features_sample = train_features_array[i_batch * num_epoch : (i_batch+1) * num_epoch] 
        _, loss_curr = sess.run(
            [fcn_solver, fcn_loss],
            feed_dict={X: train_sample, Y_features: train_features_sample, Y_adj: train_adj_sample}
            )

    if it % 10 == 0:
        print('Iter: {}; loss: {:.4}'
              .format(it, loss_curr))

print("Training finished.")

示意图

这是我对TensorFlow实现的理解。如有必要，请评论修改它。特别是，我不确定如何在示意图中表示置换矩阵的置换和​​乘法。

PyTorch

这是我的生成器的PyTorch实现。我希望能收到一些有关模型实施的反馈意见。

from torch import nn
import torch


class decoder(nn.Module):
    def __init__(self):
        super(decoder, self).__init__()

        self.n_max_atom = 9
        self.n_type_bond = 4
        self.n_feature_atom = 5

        scat_dim = 15*9
        h3_dim = 128
        h31_dim = 256
        h32_dim = 512
        h33_dim = self.n_max_atom * self.n_feature_atom * 15

        self.dec_fc30 = nn.Linear(scat_dim, h3_dim)
        self.dec_fc31 = nn.Linear(h3_dim, h31_dim)
        self.dec_fc32 = nn.Linear(h31_dim, h32_dim)
        self.dec_fc33 = nn.Linear(h32_dim, h33_dim)

        self.dec_fc40 = nn.Linear(scat_dim, h3_dim)
        self.dec_fc41 = nn.Linear(h3_dim, h31_dim)
        self.dec_fc42 = nn.Linear(h31_dim, h32_dim)
        self.dec_fc43 = nn.Linear(h32_dim, h33_dim)

        self.leaky = nn.LeakyReLU(0.2, inplace=False)

    def forward(self, z):

        h30 = self.leaky(self.dec_fc30(z))
        h31 = self.leaky(self.dec_fc31(h30))
        h32 = self.leaky(self.dec_fc32(h31))
        h33 = self.leaky(self.dec_fc33(h32))
        h33 = h33.view(-1, self.n_max_atom, self.n_feature_atom)

        h40 = self.leaky(self.dec_fc30(z))
        h41 = self.leaky(self.dec_fc31(h40))
        h42 = self.leaky(self.dec_fc32(h41))
        h43 = self.leaky(self.dec_fc33(h42))
        h43 = h43.view(-1, self.n_max_atom, self.n_type_bond, 15)
        h44 = self.leaky(torch.matmul(h43.permute(0, 2, 1, 3), h43.permute(0, 2, 3, 1)))
        h44 = h44.permute(0,2,3,1)

        return h33, h44

火车

要训练模型，我提供了包含z，f和W here的训练数据集。您可以使用以下代码加载数据。

# -- load dataset
with open('./training.data', 'rb') as f:
    feature_final = pickle.load(f)
    train_adj = pickle.load(f)
    train_features = pickle.load(f)

我将对如何为PyTorch代码实现损失功能提供一些指导。谢谢！