torch.nn.Sequential and torch.nn.Module

nn.Sequential:

seq_net	= nn.Sequential(
    nn.Linear(2,4), # define your network
    nn.Tanh(),
    nn.Linear(4,1)
)

param = seq_net.parameters() # get all parameters
optim =	torch.optim.SGD(param, 1.)

dataloader = torch.utils.data.DataLoader(
        dataset, # Inherited from torch.utils.data.Dataset
        batch_size=1,
        num_workers=1
)

criterion =	nn.BCEWithLogitsLoss()

for i,data in data_loader:
    out	= seq_net(data.x)
    loss = criterion(out, data.y)
    optim.zero_grad()
    loss.backward()
    optim.step()
    if	i % 1000 ==	0:
        print('epoch:{}, loss:{}'.format(i, loss.data[0])

nn.Module:

class net(nn.Module):
    def __init__(self, input_num, output_num):
        super(net, self).__init__()
        self.layer1 = nn.Linear(input_num, output_num)
        self.layer2 = nn.Tanh()
        self.layer3 = nn.Sequential(
            ......
        )
        ...... # define your own network layer

    def forward(self, x):
        temp1 = self.layer1(x)
        temp2 = self.layer2(temp1)
        temp3 = self.layer3(temp2)
        x = temp1 + temp2+temp3
        ...... # do whatever your want to do
        return x

jtnet = net()
param = jtnet.parameters() # get all parameters
optim =	torch.optim.SGD(param, 1.)

dataloader = torch.utils.data.DataLoader(
        dataset, # Inherited from torch.utils.data.Dataset
        batch_size=1,
        num_workers=1
)

criterion =	nn.BCEWithLogitsLoss()

for i,data in data_loader:
    out	= jtnet(data.x)
    loss = criterion(out, data.y)
    optim.zero_grad()
    loss.backward()
    optim.step()
    if	i % 1000 ==	0:
        print('epoch:{}, loss:{}'.format(i, loss.data[0]))

torch.save(jtnet.state_dict(),	'module_net.pth') # save the model

你可能感兴趣的:(人工智能)