调教神经网络时，python 和 pytorch的骚操作

用python脚本查看显卡资源

import torch

if torch.cuda.is_available():
    device_count = torch.cuda.device_count()
    print(f"Number of available GPUs: {device_count}")
    
    for i in range(device_count):
        gpu_name = torch.cuda.get_device_name(i)
        gpu_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
        print(f"GPU {i}: {gpu_name}, Memory: {gpu_memory:.2f} GB")
else:
    print("No GPUs available")

pytorch中想使用很硬的阶跃函数

import torch.nn as nn
class StepFunction(nn.Module):
    def __init__(self, edge) -> None:
        super().__init__()
        self.edge = edge
    def forward(self, x):
        return th.where(x < self.edge, th.tensor(0.0, device=x.device), th.tensor(1.0, device=x.device))