PyTorch使用

import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
import numpy as np

#network init.
class pNet(nn.Module):
    def __init__(self):
        super(pNet, self).__init__()
        self.l1 = nn.Linear(30, 60)
        self.a1 = nn.Sigmoid()
        self.l2 = nn.Linear(60, 2)
        self.a2 = nn.ReLU()
        self.l3 = nn.Softmax(dim=1)
    def forward(self, x):
        x = self.l1(x)
        x = self.a1(x)
        x = self.l2(x)
        x = self.a2(x)
        x = self.l3(x)
        return x
net = pNet()

#data load.
from sklearn.datasets import load_breast_cancer
breast_cancer = load_breast_cancer()
x_train = breast_cancer.data
y_train = breast_cancer.target

#training network.
epochs = 500
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.005)  # PyTorch suit to tiny learning rate
for epoch in range(epochs):
    optimizer.zero_grad()
    y_pred = net(x_train)
    loss = criterion(y_pred, y_train)
    loss.backward()
    optimizer.step()

#save model.
from aiworks_plugins.tf_plugins import save_torch_model_to_hdfs
save_torch_model_to_hdfs(net,torch.Tensor(1,30))