Python 中的 with 语句用法和 Pytorch 中的 with torch.no_grad() 解析

Python 中的 with 语句适用于对资源进行访问的场合，确保不管使用过程中是否发生异常都会执行必要的“清理”操作（异常处理），释放资源，比如文件使用后自动关闭／线程中锁的自动获取和释放等。例如下面是文件读取的三种写法：

# 文件读取
  
# 1) 无异常处理
file = open('file_path', 'w')
file.write('hello world !')
file.close()
  
# 2) 用 try 进行异常处理
file = open('file_path', 'w')
try:
    file.write('hello world')
except:
	print('Error !')
finally:
    file.close()

# 3) 用 with 语句进行异常处理
with open('file_path', 'w') as file:
    file.write('hello world !')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

第一种写法，可能出现两个问题：文件读取发生异常，但没有进行任何处理；或者忘记关闭文件句柄。

第二种写法就确保了文件可以正常关闭，同时如果读取发生异常也不会影响后面的代码，但缺点是比较冗长。

第三种写法就是 with 语句，同时自动完成了异常处理和文件关闭的操作。实际上 with 语句识别的是对象的 __enter__() 方法和 __exit__() 方法，只要对象有这两个方法就可以对它使用 with。而 with 相当于是先 try 一下对象的 __enter__() 方法，能调用就调用，并且调用返回的结果赋值给 as 后面的变量；然后执行 with 语句块中的代码段；执行完之后，最后 finally 要调用对象的 __exit__() 方法，如下我们自定义了一个 MessageWriter 对象：

class MessageWriter(object):
    def __init__(self, file_name):
        self.file_name = file_name
      
    def __enter__(self):
        self.file = open(self.file_name, 'w')
        return self.file
  
    def __exit__(self):
        self.file.close()
  
with MessageWriter('my_file.txt') as xfile:
    xfile.write('hello world')
1
2
3
4
5
6
7
8
9
10
11
12
13

在 Pytorch 中，tensor 有一个 requires_grad 参数，如果设置为True，这个 tensor 的梯度就会被记录，则反向传播时，该 tensor 就会自动求导。若一个节点的 requires_grad 被设置为True，那么所有依赖它的节点 requires_grad都为 True。

显然，当处于测试或者推理阶段时，我们不需要反向传播，因此也不希望内存被大量 tensor 的梯度所占用，为了节约内存，我们应该把所有 tensor 的 requires_grad 都设置为 False，with torch.no_grad() 完成的正是这个工作。

打开 torch.no_grad() 的源代码就可以看到它的 __enter__() 方法和 __exit__() 方法，如下图所示：

with torch.no_grad() 实际上就是把当前 grad 是否开启用 self.prev 记录下来，然后关闭 grad，再执行 with 语句块的代码，全部执行完之后根据 self.prev 把 grad 的状态还原回去。