三、Pytorch中getattr与self._parameters的速度对比(Python中的AI对比实验)

完整代码如下:

# -*- coding: UTF-8 -*-

# Author: Perry
# @Create Time: 2020-04-07 13:53

import timeit
import numpy as np

import torch
import torch.nn as nn


class A(nn.Module):
    def __init__(self):
        super(A, self).__init__()

        tmp = nn.Parameter(torch.Tensor(100, 100))
        self.register_parameter("node", tmp)

    def go(self):
        # _ = self._parameters["node"]
        _ = getattr(self, "node")



class B(nn.Module):
    def __init__(self):
        super(B, self).__init__()

        tmp = nn.Parameter(torch.Tensor(100, 100))
        self.register_parameter("node", tmp)

    def go(self):
        # _ = getattr(self, "node")
        _ = self._parameters["node"]


if __name__ == '__main__':
    A_exp = A()
    B_exp = B()

    repeat = 10
    number = 100000

    A_time = timeit.repeat('A_exp.go()', 'from __main__ import A_exp', repeat=repeat, number=number)
    A_time = np.mean(A_time)
    print("A_time: ", A_time)

    B_time = timeit.repeat('B_exp.go()', 'from __main__ import B_exp', repeat=repeat, number=number)
    B_time = np.mean(B_time)
    print("B_time: ", B_time)

输出结果:

A_time:  0.06433277329999998
B_time:  0.01641109059999999

结论:使用 self._parameters 字典的速度更快,但是差距不大。