Pytorch自定义网络结构+读取自己数据+自定义Loss 全过程代码示例

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/yutingzhaomeng/article/details/80454807

1、限定使用GPU的序号

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
os.system('echo $CUDA_VISIBLE_DEVICES')

2、导入相关头文件

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torch.optim as optim
from torch.autograd import Variable
import numpy as np
from Encoding import load_feature

3、自定义网络

class TransientModel(nn.Module):
    def __init__(self):
        super(TransientModel,self).__init__()
        self.conv1 = nn.Conv2d(16, 8, kernel_size=1)
        self.conv2 = nn.Conv2d(8, 4, kernel_size=1)
        self.conv3 = nn.Conv2d(4, 2, kernel_size=1)
        self.conv4 = nn.Conv2d(2, 1, kernel_size=1)
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        return x

4、自定义损失函数Loss

class MyLoss(nn.Module):
    def __init__(self):
        super(MyLoss, self).__init__()
        print '1'
    def forward(self, pred, truth):
        truth = torch.mean(truth,1)
        truth = truth.view(-1,2048)
        pred  = pred.view(-1,2048)
        return  torch.mean(torch.mean((pred-truth)**2,1),0)

5、自定义数据读取

class MyTrainData(data.Dataset):
    def __init__(self):
        self.video_path = '/data/FrameFeature/Penn/'
        self.video_file = '/data/FrameFeature/Penn_train.txt'
        fp = open(self.video_file,'r')
        lines = fp.readlines()
        fp.close()
        self.video_name = []
        for line in lines:
            self.video_name.append(line.strip().split(' ')[0])
    def __len__(self):
        return len(self.video_name)
    def __getitem__(self, index):
        data = load_feature(os.path.join(self.video_path,self.video_name[index]))
        data = np.expand_dims(data,2)
        return data

6、定义Train函数

def train(model, train_loader, myloss, optimizer, epoch):
    model.train()
    for batch_idx, train_data in enumerate(train_loader):
        train_data = Variable(train_data).cuda()
        optimizer.zero_grad()
        output = model(train_data)
        loss = myloss(output, train_data)
        loss.backward()
        optimizer.step()
        if batch_idx%100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tloss: {:.6f}'.format(
                epoch, batch_idx*len(train_data), len(train_loader.dataset),
                100.*batch_idx/len(train_loader), loss.data.cpu().numpy()[0]))

7、开始训练

if __name__=='__main__':
    # main()
    model = TransientModel().cuda()
    myloss= MyLoss()


    train_data = MyTrainData()
    train_loader = data.DataLoader(train_data,batch_size=1,shuffle=True,num_workers=1)


    optimizer = optim.SGD(model.parameters(),lr=0.001)


    for epoch in range(10):
        train(model, train_loader, myloss, optimizer, epoch)

8、结果展示

9、完整代码

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
os.system('echo $CUDA_VISIBLE_DEVICES')

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torch.optim as optim
from torch.autograd import Variable
import numpy as np
from Encoding import load_feature

class TransientModel(nn.Module):
    def __init__(self):
        super(TransientModel,self).__init__()
        self.conv1 = nn.Conv2d(16, 8, kernel_size=1)
        self.conv2 = nn.Conv2d(8, 4, kernel_size=1)
        self.conv3 = nn.Conv2d(4, 2, kernel_size=1)
        self.conv4 = nn.Conv2d(2, 1, kernel_size=1)
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        return x

class MyLoss(nn.Module):
    def __init__(self):
        super(MyLoss, self).__init__()
        print '1'
    def forward(self, pred, truth):
        truth = torch.mean(truth,1)
        truth = truth.view(-1,2048)
        pred  = pred.view(-1,2048)
        return  torch.mean(torch.mean((pred-truth)**2,1),0)

class MyTrainData(data.Dataset):
    def __init__(self):
        self.video_path = '/data/FrameFeature/Penn/'
        self.video_file = '/data/FrameFeature/Penn_train.txt'
        fp = open(self.video_file,'r')
        lines = fp.readlines()
        fp.close()
        self.video_name = []
        for line in lines:
            self.video_name.append(line.strip().split(' ')[0])
    def __len__(self):
        return len(self.video_name)
    def __getitem__(self, index):
        data = load_feature(os.path.join(self.video_path,self.video_name[index]))
        data = np.expand_dims(data,2)
        return data

def train(model, train_loader, myloss, optimizer, epoch):
    model.train()
    for batch_idx, train_data in enumerate(train_loader):
        train_data = Variable(train_data).cuda()
        optimizer.zero_grad()
        output = model(train_data)
        loss = myloss(output, train_data)
        loss.backward()
        optimizer.step()
        if batch_idx%100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tloss: {:.6f}'.format(
                epoch, batch_idx*len(train_data), len(train_loader.dataset),
                100.*batch_idx/len(train_loader), loss.data.cpu().numpy()[0]))

def main():
    model = TransientModel().cuda()
    myloss= MyLoss()

    train_data = MyTrainData()
    train_loader = data.DataLoader(train_data,batch_size=1,shuffle=True,num_workers=1)

    optimizer = optim.SGD(model.parameters(),lr=0.001)

    for epoch in range(10):
        train(model, train_loader, myloss, optimizer, epoch)

if __name__=='__main__':
    main()