Pytorch 神经网络—自定义数据集上实现

LZDCQU

于 2018-09-11 21:58:28 发布

阅读量5.6k

点赞数 1

分类专栏：机器学习

本文链接：https://blog.youkuaiyun.com/qq_21905401/article/details/82627402

版权

第一步、导入需要的包

import os
import scipy.io as sio
import numpy as np
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from torch.autograd import Variable

batchSize = 128 # batchsize的大小
niter = 10      # epoch的最大值

第二步、构建神经网络

设神经网络为如上图所示，输入层4个神经元，两层隐含层各4个神经元，输出层一个神经。每一层网络所做的都是线性变换，即y=W×X+b；代码实现如下：

class Neuralnetwork(nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Neuralnetwork, self).__init__()
        self.layer1 = nn.Linear(in_dim, n_hidden_1)
        self.l