def setup_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
下面的代码开启 benchmark 模式。该模式针对网络的输入大小不变的情况下是好的。
cudnn 将会花一些时间为特定的配置找到理想的算法集合(比如卷积的具体实现)。
这通常使得在运行的时候更快。
但是如果你的 输入大小在每个 iteration 都会变化,那么cudnn将会在每次 输入大小变化的时候 benchmark,有可能会导致更差的 运行时表现。
这取决于任务,如果输入大小变化很大,那么它可能会延长运行时长,反之,它应该会加速。
torch.backends.cudnn.benchmark = True