PyTorch转TensorFlow：如何实现张量拼接及代码转换？

安全阿岳

2026-5-25

Hey there! Let's break down how to convert your PyTorch code to TensorFlow step by step, since you're just getting started with TF. Here's the equivalent code first, then we'll walk through each part to make it clear:

import tensorflow as tf

# Split dataset into train/test sets and convert to TensorFlow tensors
X_train = tf.convert_to_tensor(X[:-30], dtype=tf.float32)
Y_train = tf.convert_to_tensor(Y[:-30], dtype=tf.float32)
X_test = tf.convert_to_tensor(X[-30:], dtype=tf.float32)
Y_test = tf.convert_to_tensor(Y[-30:], dtype=tf.float32)

# Flatten tensors: keep the sample count dimension, flatten the rest into one dimension
X_train = tf.reshape(X_train, (tf.shape(X_train)[0], -1))
Y_train = tf.reshape(Y_train, (tf.shape(Y_train)[0], -1))

# Concatenate along the second axis (equivalent to PyTorch's dim=1)
data = tf.concat([X_train, Y_train], axis=1)

# Print the tensor type
print(type(data))

Key differences and explanations compared to your PyTorch code:

Tensor conversion: In PyTorch you used Variable(torch.Tensor(...)), but in TensorFlow 2.x, we don't need a separate Variable wrapper for raw data tensors (that's only for trainable parameters). tf.convert_to_tensor() directly turns your input data into a TensorFlow tensor, and dtype=tf.float32 matches PyTorch's default float tensor type.
Flattening: PyTorch's view(X_train.size(0), -1) maps directly to TensorFlow's tf.reshape(). tf.shape(X_train)[0] gets the number of samples, and -1 tells TensorFlow to automatically calculate the total length of the remaining dimensions—exactly the same flattening behavior. If you prefer using Keras layers, you could also do tf.keras.layers.Flatten()(X_train), but tf.reshape() is more straightforward for raw tensor operations.
Concatenation: torch.cat(..., dim=1) is tf.concat(..., axis=1) in TensorFlow. The parameter name changes, but the functionality is identical: we're joining the two tensors along the column dimension.
Type check: When you print type(data), you'll get <class 'tensorflow.python.framework.ops.EagerTensor'>—this is TensorFlow 2.x's default eager execution tensor type, analogous to PyTorch's Tensor class.