第二大脑 / Knowledge Bank

with-context-manager

7分钟阅读

with 与上下文管理器 (Context Manager)

1. 定义

with 是 Python 的上下文管理器语法：把”进入前做准备、退出后自动收尾”这件事打包起来。最典型的作用是自动释放资源——不管代码正常结束还是中途报错，with 块结束时都会自动执行清理。

with open("data.txt") as f:      # 进入：打开文件
    content = f.read()
# 退出：自动关闭文件，即使上面报错也会关

一句话：with = “用完自动关”。 省得你手动写 close()，还能保证异常时也不漏关。

2. 解决了什么问题

没有 with 时要手动收尾，还得用 try/finally 防异常漏掉：

# 笨办法：手动 + try/finally
f = open("data.txt")
try:
    content = f.read()
finally:
    f.close()            # 必须保证关，否则文件句柄泄露
 
# with：一行搞定，自动等价于上面的 try/finally
with open("data.txt") as f:
    content = f.read()

with 把”进入 → 使用 → 无论如何都收尾”这套模式封装了，代码更短、更安全。

3. 工作原理（简单了解）

一个对象只要实现了 __enter__（进入时调用）和 __exit__（退出时调用，含异常时）两个方法，就能用于 with：

class MyResource:
    def __enter__(self):
        print("进入：准备资源")
        return self                 # as 后面拿到的就是这个返回值
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("退出：清理资源")      # 正常/异常都会执行
 
with MyResource() as r:
    print("使用中")
# 输出：进入 → 使用中 → 退出

as x 拿到的是 __enter__ 的返回值。__exit__ 的三个参数是异常信息（没异常时都是 None）。

4. 与 PyTorch 无关的常见场景

这些是最基础、最通用的用法：

① 文件读写（最经典）

with open("out.txt", "w") as f:
    f.write("hello")
# 自动 flush + close

② 多个资源一起管理

with open("in.txt") as fin, open("out.txt", "w") as fout:
    fout.write(fin.read())

③ 数据库连接 / 网络请求

import sqlite3
with sqlite3.connect("db.sqlite") as conn:
    conn.execute("INSERT INTO t VALUES (1)")
# 自动提交/回滚
 
import requests
with requests.get("https://...", stream=True) as resp:
    data = resp.content

④ 线程锁（并发时自动加锁/解锁）

import threading
lock = threading.Lock()
with lock:                 # 进入自动 acquire，退出自动 release
    shared_counter += 1

⑤ 临时改变状态（用 contextlib 自己造）

from contextlib import contextmanager
 
@contextmanager
def timer():
    import time
    start = time.time()
    yield                          # yield 之前 = __enter__，之后 = __exit__
    print(f"耗时 {time.time()-start:.2f}s")
 
with timer():
    do_something()                 # 自动计时并打印

@contextmanager + yield 是”快速造一个上下文管理器”的常用手法：yield 前是进入逻辑，yield 后是退出逻辑。

5. 与 PyTorch 有关的常见场景

PyTorch 大量用 with 来临时切换某种全局/计算状态，块结束自动还原：

① torch.no_grad()：推理时关闭梯度追踪（最高频！）

import torch
 
model.eval()
with torch.no_grad():              # 这个块内不构建计算图、不追踪梯度
    for x, y in test_loader:
        pred = model(x)            # 省显存、跑得快
        ...
# 出块后梯度追踪自动恢复

为什么用它：推理/评估时不需要反向传播，关掉 autograd 能省显存、加速。这是 with 在 PyTorch 里最常见的用法。见 torch.md 第 6、10 节。

② torch.inference_mode()：比 no_grad 更彻底的推理模式

with torch.inference_mode():       # 比 no_grad 更省，推理专用（PyTorch 1.9+）
    pred = model(x)

和 no_grad 类似但优化更激进，纯推理场景推荐。

③ torch.autocast：混合精度训练（AMP）

with torch.autocast(device_type="cuda", dtype=torch.float16):
    output = model(input)          # 块内自动用半精度算，省显存、提速
    loss = loss_fn(output, target)

④ torch.cuda.device：临时指定用哪块 GPU

with torch.cuda.device(1):         # 块内默认用 1 号 GPU
    x = torch.randn(3, 3).cuda()

⑤ 对比：训练 vs 推理的写法差异

# 训练：需要梯度，不用 no_grad
for x, y in train_loader:
    pred = model(x)
    loss = loss_fn(pred, y)
    optimizer.zero_grad()
    loss.backward()                # 要反向传播，必须有计算图
    optimizer.step()
 
# 推理/评估：不需要梯度，用 no_grad 包起来
with torch.no_grad():
    for x, y in test_loader:
        pred = model(x)            # 只前向，省资源

共同点：这些 PyTorch 的 with 都是**“进入时切换某个全局状态，退出时自动还原”**——和文件的”打开/关闭”是同一个思想，只是管理的不是文件句柄，而是”要不要追踪梯度 / 用什么精度 / 用哪块卡”。

6. 常见坑

  • 该用 no_grad 却忘了：推理时不加 torch.no_grad()，会白白建计算图、爆显存、变慢（PyTorch 头号坑之一）
  • 以为 no_grad 等于 eval：两者不同！model.eval() 切换 Dropout/BatchNorm 行为，torch.no_grad() 关梯度追踪——推理时通常两个都要
  • 在 no_grad 块里想 backward：块内张量没有梯度图，loss.backward() 会报错；训练不能包在 no_grad 里
  • 文件没用 with 手动 open 忘了 close：文件句柄泄露；能用 with 就别手动 open
  • with 块外还想用块内资源：出了 with 文件已关闭，再读会报错
  • @contextmanager 里异常处理：如果希望异常时也执行退出逻辑，要用 try/finally 包住 yield

7. 延伸阅读 / 关联概念

  • PyTorch — no_grad / eval / autocast 的完整用法；见 torch.md
  • CUDA — torch.cuda.device、显存管理；见 cuda.md
  • try/finally — with 底层等价的异常安全机制
  • contextlib — @contextmanager、ExitStack、suppress 等造上下文管理器的工具
  • __enter__ / __exit__ — 自定义上下文管理器的两个魔术方法
  • 官方文档：with statement ｜ contextlib

关系图谱