49a2ea975624090307c652e91258e5b6f02cda41,examples/tetris.py,,main,#,180

Before Change

    network = SE3Net().cuda()
    train(network, trainset, N_epochs=N_epochs)
    se3_test_accs = []
    for _ in range(N_test):
        testset = get_volumes(rotate=True)
        acc = test(network, testset)
        se3_test_accs.append(acc)

    // network = CNN().cuda()
    // train(network, trainset, N_epochs=N_epochs)
    // cnn_test_accs = []
    // for _ in range(N_test):
    //     testset = get_volumes(rotate90=True)
    //     acc = test(network, testset)
    //     cnn_test_accs.append(acc)

    print("avg test acc SE3: {}".format(np.mean(se3_test_accs)))
    // print("avg test acc CNN: {}".format(np.mean(cnn_test_accs)))
    // N_classes = len(testset[1])
    // print("random guessing accuracy: {}".format(1 / N_classes))

After Change

def main():
    torch.backends.cudnn.benchmark = True

    for smooth_stride in [True, False]:
        for Model in [SE3Net, CNN]:
            for _rep in range(5):
                network = Model(smooth_stride).cuda()
                acc = experiment(network)
                print("smooth_stride={} model={} acc= {}".format(smooth_stride, Model, acc))


main()

Instances