tensorflow和pytorch的对应函数。

tf.math.square ---------- torch.pow

tf.reduce_sum() ---------torch.sum

tf.boolean_mask ----------torch.masked_select（必须是bool值）

tf.reduce_mean() ---------torch.mean

tf.cast() ---------- x.type() or x.to()

tf.stack() ------------ torch.stack()

在用pytorch训练模型时，通常会在遍历epochs的过程中依次用到optimizer.zero_grad(),loss.backward()和optimizer.step()三个函数，如下所示：

for epoch in range(1, epochs):
    for i, (inputs, labels) in enumerate(train_loader):
        output= model(inputs)
        loss = criterion(output, labels)
        
        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()