📝 Common PyTorch Code Snippets

Tue, 21 Mar 2023 00:00:00 +0000

Notes inspired by d2l’s PyTorch content

Device Switching Link to heading

Try using GPU:

def try_gpu(i=0):
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():
    """Find all available GPUs"""
    devices = [torch.device(f'cuda:{i}')
               for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

Accumulator Link to heading

class Accumulator:
    """Used for accumulating values"""
    def __init__(self, n):
        self.data = [0.0] * n

    def add(self, *args):
        """Add a list of values to the existing data"""
        self.data = [a + float(b) for a, b in zip(self.data, args)]

    def reset(self):
        self.data = [0.0] * len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

Accuracy Link to heading

For training accuracy evaluation, you can use the following code

Snippets on ChengAo Shen

📝 Common PyTorch Code Snippets

Device Switching Link to heading

Accumulator Link to heading

Accuracy Link to heading