XI. Advanced Features & Utilities (高级特性与实用工具)#

1
optimizer.zero_grad()
2
loss.backward()
3
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
4
optimizer.step()

1
from torch.nn.utils import weight_norm, remove_weight_norm
2
wn_conv = weight_norm(nn.Conv1d(64, 64, 3, padding=1))
3
remove_weight_norm(wn_conv)  # merge before deployment

1
x = torch.rand(1, 64, 28, 28)
2
up = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
3
print(up.shape)  # [1, 64, 56, 56]

1
theta = torch.eye(2, 3, dtype=torch.float).unsqueeze(0)
2
grid = F.affine_grid(theta, x.size())
3
out = F.grid_sample(x, grid, mode='bilinear')

1
c = torch.einsum('ij,jk->ik', a, b)  # matrix multiply
2

3
# Attention scores: Q:[B,H,L,D], K:[B,H,L,D]
4
scores = torch.einsum('bhld,bhmd->bhlm', Q, K)

1
from torch.profiler import profile, ProfilerActivity
2
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
3
    model(x)
4
print(prof.key_averages().table(sort_by='cuda_time_total'))

1
def init_weights(m):
2
    if isinstance(m, nn.Linear):
3
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
4
        nn.init.zeros_(m.bias)
5

6
model.apply(init_weights)

1
loss = criterion(output, label)
2
loss_val = loss.item()        # detaches from graph
3
print(f'Loss: {loss_val:.4f}')
4

5
# WRONG: total_loss += loss  ← graph grows → OOM
6
# CORRECT:
7
total_loss += loss.item()

1
x = torch.rand(3, requires_grad=True)
2
y = x.clone()           # gradient can still propagate
3
z = x.detach().clone()  # gradient detached
4

5
buf = torch.empty(3)
6
buf.copy_(x)            # in-place copy

1
x = torch.tensor([-1., 2., -3., 4.])
2
y = torch.zeros_like(x)
3
out = torch.where(x > 0, x, y)  # tensor([0., 2., 0., 4.])  — manual ReLU

1
logits = torch.rand(4, 10)
2
targets = torch.tensor([3, 7, 1, 5]).unsqueeze(1)  # [4, 1]
3
scores = logits.gather(dim=1, index=targets)         # [4, 1]

1
y = torch.zeros(4, 5)
2
labels = torch.tensor([[2], [0], [4], [1]])
3
y.scatter_(dim=1, index=labels, value=1.0)  # one-hot encoding

1
L = 5
2
mask = torch.triu(torch.ones(L, L), diagonal=1).bool()
3
scores = torch.rand(L, L)
4
scores = scores.masked_fill(mask, float('-inf'))  # causal mask

1
from torch.nn.utils.rnn import pad_sequence
2
seqs = [torch.tensor([1, 2, 3]), torch.tensor([4, 5]), torch.tensor([6])]
3
padded = pad_sequence(seqs, batch_first=True)  # [3, 3]

1
enc_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, dim_feedforward=2048, batch_first=True)
2
encoder = nn.TransformerEncoder(enc_layer, num_layers=6)
3
out = encoder(src, src_key_padding_mask=m)

1
x = torch.rand(4, 5, 6)
2
y = x.permute(2, 0, 1)         # non-contiguous
3
print(y.is_contiguous())        # False
4
z = y.contiguous()
5
w = z.view(6, -1)               # safe to view now

1
x = torch.tensor([1, 2, 3])
2
f = x.float()   # int → float32
3
h = x.half()    # float32 → float16
4
l = x.long()    # → int64
5
y = x.to(dtype=torch.float32)  # recommended