trorch.nn使用

重要的一个类：module

import torch.nn as nn
import torch.nn.functional as F

#nn.Module代表继承nn.Module
class Model(nn.Module):
    #初始化函数
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)
    #传播，进行一次卷积+一次非线性激活
    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

举例：

import torch
from torch import nn

# 重写类，快捷键Alt+insert
class module(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input):
        output = input+1
        return output
# 新建对象
module1 = module()
x = torch.tensor([1,0])
output =module1(x)
print(output)

 初始化与一个tensor的数据

torch.tensor([[1., -1.], [1., -1.]])
torch.tensor(np.array([[1, 2, 3], [4, 5, 6]]))


#或者
torch.zeros([2, 4], dtype=torch.int32)
cuda0 = torch.device('cuda:0')
torch.ones([2, 4], dtype=torch.float64, device=cuda0)