Question

我刚开始使用PyTorch，我想解决一些玩具问题。在以下情况下，我注意到模型一次训练并发布一批预测需要花费多少时间。

这是PyTorch的实现。在GPU上，我的计算机大约需要17秒。在CPU上使用相同的型号大约需要11秒钟。

class LR(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(2, 20)
        self.linear2 = torch.nn.Linear(20, 1)

    def forward(self, x):
        x = torch.nn.functional.relu(self.linear1(x))
        x = torch.nn.functional.relu(self.linear2(x))
        return x


def fit_torch(df_train, df_test):
    sampler_tr = torch.utils.data.SubsetRandomSampler(df_train.index)
    train = torch.utils.data.DataLoader(
        torch.tensor(df_train.values, dtype=torch.float),
        batch_size=batch_size, sampler=sampler_tr)

    sampler_te = torch.utils.data.SubsetRandomSampler(df_test.index)
    test = torch.utils.data.DataLoader(
        torch.tensor(df_test.values, dtype=torch.float),
        batch_size=batch_size, sampler=sampler_te)

    model = LR()
    model = model.to(device)

    loss = torch.nn.MSELoss()
    optim = torch.optim.Adam(model.parameters(), lr=0.001)

    model.train()
    for _ in range(1000):
        for train_data in train:
            train_data = train_data.to(device)

            x_train = train_data[:, :2]
            y_train = train_data[:, 2]

            optim.zero_grad()

            pred = model(x_train)
            loss_val = loss(pred.squeeze(), y_train)

            loss_val.backward()
            optim.step()

    model.eval()
    with torch.no_grad():
        for test_data in test:
            test_data = test_data.to(device)

            pred = model(test_data[:, :2].float())
            break

这是keras实现。运行大约需要9秒钟。

def fit_tf(df_train, df_test):
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Dense(20, activation='relu'))
    model.add(tf.keras.layers.Dense(1, activation='relu'))

    model.compile(loss='mse', optimizer='adam')
    model.fit(
        df_train.values[:, :2],
        df_train.values[:, 2],
        batch_size=batch_size, epochs=1000, verbose=0)

    model.predict(df_test.iloc[:batch_size].values[:, :2])

数据集和主要功能。

device = torch.device('cuda:0')
scaler = MinMaxScaler()

batch_size = 64

def create_dataset():
    dataset = []    
    random_x = np.random.randint(10, 1000, 1000)
    random_y = np.random.randint(10, 1000, 1000)

    for x, y in zip(random_x, random_y):
        dataset.append((x, y, 4 * x + 3 * y + 10))

    np.random.shuffle(dataset)
    df = pd.DataFrame(dataset)
    df = pd.DataFrame(scaler.fit_transform(df))

    return df

def __main__():
    df = create_dataset()
    df_train, df_test = train_test_split(df)

    start_time = time.time()
    fit_tf(df_train.reset_index(drop=True), df_test.reset_index(drop=True))
    print(time.time() - start_time)

Answer 1

PyTorch默认情况下使用动态计算图，当您开始开发神经网络时，它会更加灵活，因为它将提供更直接的调试消息。相反，TensorFlow将生成一个静态计算图，这就是为什么您需要在使用模型之前对其进行编译。编译器可以优化模型，但代价是神经网络变得难以调试。这可能会导致两个框架之间的性能存在细微的差异，但这并不重要。

由于您的网络非常小，因此在GPU内存和CPU内存之间复制网络以及启动CUDA子系统的开销超过了GPU带来的好处。如果尝试使用诸如AlexNet，ResNet甚至GoogLeNet之类的更复杂的神经网络，则好处将更加明显。

PyTorch（GPU）慢于CPU慢于keras

1 个答案: