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