# version
torch.__version__               # PyTorch version
torch.version.cuda              # Corresponding CUDA version
torch.backends.cudnn.version()  # Corresponding cuDNN version
torch.cuda.get_device_name(0)   # GPU type

# seed
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

# GPU
torch.cuda.is_available()
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
torch.cuda.empty_cache()        # clear GPU cache

device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)

tensor = torch.randn(2, 3)

# tensor info
tensor.type()   # Data type
tensor.size()   # Shape of the tensor
tensor.shape    # Shape of the tensor
tensor.dim()    # Number of dimensions

# type convertions
tensor = tensor.cuda()
tensor = tensor.cpu()
tensor = tensor.float()
tensor = tensor.long()
torch.set_default_tensor_type(torch.FloatTensor)   # Set default tensor type

# torch.Tensor <=> np.ndarray
ndarray = tensor.cpu().numpy()
tensor = torch.from_numpy(ndarray).float()
tensor = torch.from_numpy(ndarray.copy()).float()  # If ndarray has negative stride

# 从只包含一个元素的张量中提取值
value = tensor.item()

# torch.Tensor <=> PIL.Image
# PyTorch中的张量默认采用N×D×H×W的顺序，并且数据范围在[0, 1]，需要进行转置和规范化。
image = PIL.Image.fromarray(torch.clamp(tensor * 255, min=0, max=255
    ).byte().permute(1, 2, 0).cpu().numpy())
image = torchvision.transforms.functional.to_pil_image(tensor)  # Equivalently way

tensor = torch.from_numpy(np.asarray(PIL.Image.open(path))
    ).permute(2, 0, 1).float() / 255
tensor = torchvision.transforms.functional.to_tensor(PIL.Image.open(path))  # Equivalently way

# np.ndarray <=> PIL.Image
image = PIL.Image.fromarray(ndarray.astypde(np.uint8))
ndarray = np.asarray(PIL.Image.open(path))

# reshape
tensor.reshape(shape)
tensor.permute(0,2,1) # swap axes
tensor.flatten() # 展成 1D 向量
tensor.squeeze() # 去掉维数为 1 的的维度
tensor.unsqueeze(dim=0) # 增加维度

# shuffle
tensor = tensor[torch.randperm(tensor.size(0))]  # Shuffle the first dimension

# copy
# Operation                 |  New/Shared memory | Still in computation graph |
tensor.clone()            # |        New         |          Yes               |
tensor.detach()           # |      Shared        |          No                |
tensor.detach().clone()   # |        New         |          No                |

# concatenate
tensor = torch.cat(list_of_tensors, dim=0)
tensor = torch.stack(list_of_tensors, dim=0)

# one-hot
N = tensor.size(0)
one_hot = torch.zeros(N, num_classes).long()
one_hot.scatter_(dim=1,
    index=torch.unsqueeze(tensor, dim=1),
    src=torch.ones(N, num_classes).long())

# non-zero
torch.nonzero(tensor)               # Index of non-zero elements
torch.nonzero(tensor == 0)          # Index of zero elements
torch.nonzero(tensor).size(0)       # Number of non-zero elements
torch.nonzero(tensor == 0).size(0)  # Number of zero elements

# equal
torch.allclose(tensor1, tensor2)  # float tensor
torch.equal(tensor1, tensor2)     # int tensor

# expand
# Expand tensor of shape 64*512 to shape 64*512*7*7.
torch.reshape(tensor, (64, 512, 1, 1)).expand(64, 512, 7, 7)

# normalize
F.normalize(tensor, dim=1)

from torch.utils.data import TensorDataset, Dataset, DataLoader

train_ds = TensorDataset(x_train[:50000], y_train[:50000])
valid_ds = TensorDataset(x_train[50000:], y_train[50000:])
train_dl = DataLoader(train_ds, batch_size=mini_batch, shuffle=True, num_workers=4)
valid_dl = DataLoader(valid_ds, batch_size=len(valid_ds))

for xb, yb in train_dl:
    pass

# 需要实现__getitem__、__len__两个方法
class SparseDataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __getitem__(self, index):
        data = self.data[index].tocoo()
        i = torch.LongTensor(np.vstack((data.row, data.col)))
        v = torch.FloatTensor(data.data)
        data = torch.sparse.FloatTensor(i, v, torch.Size(data.shape))
        return data
    
    def __len__(self):
        return self.data.shape[0]

# convolution layer
conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True)
conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=True)

init.kaiming_normal_(layer.weight, mode='fan_out', nonlinearity='relu')
nn.init.constant_(layer.weight, val=1.0)
nn.init.constant_(layer.bias, val=0.0)
nn.init.xavier_normal_(layer.weight)
nn.init.constant_(layer.bias, val=0.0)

# Initialization with given tensor.
layer.weight = torch.nn.Parameter(tensor)

# init
super(Model, self).__init__()

# parameter
self.weight = nn.Parameter(torch.Tensor(n_classes, embed_dim))
init.xavier_uniform_(self.weight)

# parameter list
self.params = nn.ParameterList()
for i in range(3):
    self.params.append(nn.Parameter(torch.Tensor(2,3)))

idx = np.arange(n_train)
for i in range(epochs):
    np.random.shuffle(idx)
    tr = trange(n_train // batch + 1)
    for j in tr:
        batch_idx = idx[j*batch:(j+1)*batch]
        pred = model(x_train[batch_idx])
        loss = model.loss(pred, labels[batch_idx])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        tr.set_description("Epoch {}: loss: {}".format(i, loss.item()))
        tr.refresh()

# If there is one global learning rate (which is the common case).
lr = next(iter(optimizer.param_groups))['lr']

# If there are multiple learning rates for different layers.
all_lr = []
for param_group in optimizer.param_groups:
    all_lr.append(param_group['lr'])

# Reduce learning rate when validation accuarcy plateau.
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=5, verbose=True)
for t in range(0, 80):
    train(...); val(...)
    scheduler.step(val_acc)

# Cosine annealing learning rate.
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=80)
# Reduce learning rate by 10 at given epochs.
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[50, 70], gamma=0.1)
for t in range(0, 80):
    scheduler.step()
    train(...); val(...)

# Learning rate warmup by 10 epochs.
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda t: t / 10)
for t in range(0, 10):
    scheduler.step()
    train(...); val(...)

# Save checkpoint.
is_best = current_acc > best_acc
best_acc = max(best_acc, current_acc)
checkpoint = {
    'best_acc': best_acc,    
    'epoch': t + 1,
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
}
model_path = os.path.join('model', 'checkpoint.pth.tar')
torch.save(checkpoint, model_path)
if is_best:
    shutil.copy('checkpoint.pth.tar', model_path)

# Load checkpoint.
if resume:
    model_path = os.path.join('model', 'checkpoint.pth.tar')
    assert os.path.isfile(model_path)
    checkpoint = torch.load(model_path)
    best_acc = checkpoint['best_acc']
    start_epoch = checkpoint['epoch']
    model.load_state_dict(checkpoint['model'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    print('Load checkpoint at epoch %d.' % start_epoch)