PyTorch与TensorFlow在Darknet中的速度对比

前言

最近开始学习pytorch，之前我用的是tensorflow2.4，体验实在是一言难尽，现在有时间了，还是决定转到pytorch上感受一下，以下是我在分类问题中的一些总结，供大家参考。

• 上手难度方面。我推荐pytorch，虽然tensorflow从代码上看更简洁，有些功能打包的更完整，但是由于2.0版本和1.0版本差别较大，会让你在学习过程中踩很多坑，弯路真是走了不少啊。
• 参考资料方面。相信大家在网上会找不少教程，不难发现tensorflow的帖子比较少，也比较老，很多新开源的代码也以pytorch为主，相关资源找起来实在是费劲。
• 工程实践方面。据说tensorflow还是很好用的，性能更强，应用也比较广泛。pytorch则处于追赶阶段，工程应用还是无法取代前者的。

下面是我在轴承故障分类任务中感受比较深的几个点。

在数据调用方面。

tensorflow中由于轴承数据大小尚可，直接使用numpy编写后送入模型进行训练，基本是这样的：

登录后复制

trainX, trainy, valX, valy, testX, testy = base_data1

只要大小合适，拿来就能用。

而在pytorch中，如果还需要转为torch自己的格式，然后生成迭代器，再进行训练，用起来细节更多

登录后复制

data_tr, label_tr, data_val, label_val, data_te, label_te = data
data_tr, label_tr, data_val, label_val, data_te, label_te = map(
        torch.tensor, (data_tr, label_tr, data_val, label_val, data_te, label_te))
data_train = Data.TensorDataset(data_tr, label_tr)
data_val = Data.TensorDataset(data_val, label_val)
data_te = Data.TensorDataset(data_te, label_te)
data_loader_train = Data.DataLoader(dataset=data[0], batch_size=BATCH_SIZE, shuffle=True)

训练方面。

tensorflow也比较简单一些，一个model.fit搞定大部分的训练、验证、模型保存等过程，你只需要把数据集整规范就行

登录后复制

model.fit(trainX, trainy, epochs=epochs, batch_size=128, verbose=0, validation_data=(valX, valy))

pytorch中同样的内容相当于拆开了进行，你可以自定义自己的函数，但是很多地方还是需要自己操心，同样的内容，代码就不少了，下面我只写一部分

登录后复制

def train(data,model):  #在指定数据集上进行训练
    model.train()  # 作用是启用batch normalization和drop out。
    running_loss = 0.0  # 将每个batch_size的损失加到一起
    train_acc = 0
    criterion = nn.CrossEntropyLoss().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)
    data_loader_train = Data.DataLoader(dataset=data[0], batch_size=BATCH_SIZE, shuffle=True)

    for X, y in data_loader_train:
        input = X.to(device)  # 将数据拷贝到GPU上
        y = y.to(device)
        output = model.forward(input)  # 将数据前向传入训练模型
        loss = criterion(output, y)
        running_loss += loss.item()

        # 这一部分永远是固定的
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        _, pred = torch.max(output, 1)
        _, y_train = torch.max(y, 1)

        train_acc += torch.sum(pred == y_train)

    epoch_loss_train = running_loss / len(data[0])
    train_acc = train_acc / len(data[0])

    return train_acc,epoch_loss_train

def val(data,model):    #使用验证集对训练结果进行验证

    model.eval（)  # 模型验证——关闭drop out
    val_loss = 0.0
    val_acc = 0.0
    criterion = nn.CrossEntropyLoss().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)
    data_loader_val = Data.DataLoader(dataset=data[1], batch_size=BATCH_SIZE, shuffle=True)
    with torch.no_grad():  # 不用梯度了
        for X, y in data_loader_val:
            input = X.to(device)  # 将数据拷贝到GPU上
            optimizer.zero_grad()  # 梯度置零
            output = model(input)
            y = y.to(device)
            val_loss += criterion(output,y).item()
            val_acc += (output.argmax(1) == y.argmax(1)).type(torch.float).sum().item()

    epoch_loss_val = val_loss / len(data[1])
    val_acc = val_acc / len(data[1])

    return val_acc,epoch_loss_val
    
    

以上还只是训练加验证的内容。当然，从原理讲大家都是一样的，pytorch确实更有利于自己把握细节，有错误也更好找。

总而言之就是tensorflow打包的很严实，就算自己不太懂，按部就班也不影响使用，但是一旦出现问题，你还得查源代码，这时候就比较难受了，问题确实也难找一些。pytorch因为一开始就比较细吧，整个过程下来以后相对顺一些，有问题也比较好找。

免责声明：本文系网络转载或改编，未找到原创作者，版权归原作者所有。如涉及版权，请联系删