PyTorch性能优化10个技巧：效率翻倍的实战雕虫！

import torch
from torch.utils.bottleneck import bottleneck

model = YourModelDefinition()
inputs = torch.randn(batch_size, input_shape)

with bottleneck(mode='full'):
    outputs = model(inputs)
    loss = loss_function(outputs)
    loss.backward()

bottleneck.show_bottlenecks()

import torch
from torch import nn

# Start recording memory snapshot history
torch.cuda.memory._record_memory_history(max_entries=100000)

# Example model and computation
model = nn.Linear(10_000, 50_000, device="cuda")
for _ in range(3):
    inputs = torch.randn(5_000, 10_000, device="cuda")
    outputs = model(inputs)

# Dump memory history to a file and stop recording
torch.cuda.memory._dump_snapshot("profile.pkl")
torch.cuda.memory._record_memory_history(enabled=None)

# To install TorchLens, first install graphviz
# (required to generate the network visualizations)
sudo apt install graphviz

# and then install TorchLens using pip
pip install torchlens -q

class SimpleRecurrent(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(in_features=5, out_features=5)

    def forward(self, x):
        for r in range(4):
            x = self.fc(x)
            x = x + 1
            x = x * 2
        return x


simple_recurrent = SimpleRecurrent()
model_history = tl.log_forward_pass(simple_recurrent, x,
                                    layers_to_save='all',
                                    vis_opt='rolled')
print(model_history['linear_1_1:2'].tensor_contents)
# second pass of first linear layer

"""
tensor([[-2.9265, 1.3072, -0.4358, -0.9827, -3.0030],
        [-2.1188, 0.6454, 0.1414, -1.0781, -2.4770],
        [-1.8614, 1.0404, -0.7678, -0.7843, -2.4568],
        [-4.4318, 1.4188, -1.6984, -0.3384, -5.2524],
        [-1.8989, 0.8374, -0.7154, -0.6169, -2.7501],
        [-3.2618, 1.5131, 1.2455, -1.7886, -1.5585]],
       grad_fn=                                                     0>)
"""

# OK Code
tensor = torch.ones([100, 2]) 
tensor = tensor.cuda()

# Better Code
tensor = torch.ones([100, 2], device="cuda:0")

import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("mask", torch.ones(10))

    def forward(self, x):
        return x * self.mask

# Example usage
model = MyModel()
print(model.mask)  # Persistent buffer, not a parameter

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Dummy dataset
x, y = torch.randn(1000, 10), torch.randn(1000, 1)
train_loader = DataLoader(TensorDataset(x, y), 
              batch_size=32, shuffle=True)

# Simple model
model = nn.Linear(10, 1)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

accumulation_steps = 4

for epoch in range(5):
    for i, (inputs, labels) in enumerate(train_loader):
        loss = criterion(model(inputs), labels) / accumulation_steps
        loss.backward()
        if (i + 1) % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

    print(f"Epoch {epoch+1}, Loss: {loss.item()}")

print("Training completed")

import torch
from transformers import BertModel, BertTokenizer

# Load the small pre-trained BERT model
tokenizer = BertTokenizer.from_pretrained("prajjwal1/bert-small")
model = BertModel.from_pretrained("prajjwal1/bert-small")

# Example text input
text = "Hello, how are you?"

# Tokenize the input text
inputs = tokenizer(text, return_tensors="pt")

# Efficient inference
with torch.no_grad():
    output = model(**inputs)

# Detach tensor
detached_tensor = output.last_hidden_state.detach()

!pip install torchdata

from torchdata.datapipes.iter import IterableWrapper, Mapper, Filter

# Create a simple dataset
data = IterableWrapper(range(10))

# Define preprocessing steps
def square(x): return x ** 2
def is_even(x): return x % 2 == 0

# Chain operations
processed_data = data.map(square).filter(is_even)

# Use the processed data
for item in processed_data:
    print(item)

from torchft import Manager, DistributedDataParallel
from torchft import Optimizer, ProcessGroupGloo

manager = Manager(pg=ProcessGroupGloo())
model = nn.Linear(2, 3)
model = DistributedDataParallel(manager, model)
optimizer = Optimizer(manager, optim.AdamW(model.parameters()))

for epoch in range(1000):
    optimizer.zero_grad()
    output = model(torch.rand(2, 2))
    loss = output.sum()
    loss.backward()
    optimizer.step()

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

model = Model()
scripted_model = torch.jit.script(model)
scripted_model.save("model.pt")