pytorch 下自定义激活函数与多GPU训练

 一、在解决一些实际问题时,pytorch中自带的激活函数可能不能满足需求,就需要自定义一些激活函数,比如我需要一个使得输出值在0到140之间的激活层,pytorch中不含有,所以自定义:

#define a self activation function
class Act_fun(nn.Module):
    def __init__(self):
        super(Act_fun, self).__init__()
        
    def forward(self, x):
        x = F.sigmoid(x)
        x = x * 140
        return x 

         在之后的使用中,直接调用Act_op()即可,如:

self.linear2 = nn.Sequential(nn.Linear(1024,1980),Act_fun())

二、多块GPU下的同时训练

目前用的服务器中有6块1080Ti,编号是从0-5,为了全部使用,需要用到pytorch中的nn.DataParallel()函数来进行数据的并行:

    #Define Network
    model = Net(args.input_channel,args.output_channel)
    #using the multiple GPU to train the model
    if torch.cuda.device_count()>1:
        model = nn.DataParallel(model)
    model.to(device)

DataParallel()中有个参数,divice_ids,当不设置值是使用的是全部的GPU,设置值时,如divice_ids=[0,1,2,3],使用的第一块到第四块GPU。不过这种方法存在不均衡的问题,即主GPU利用率高而次GPU利用率低。

  • 0
    点赞
  • 14
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值