pytorch学习笔记(二) 其他机制

broadcast机制

pytroch中的tensor和numpy中的ndarray一样是支持broadcast的。支持原则是:

1.两个tensor都至少有一维

2.从最后一维开始,要么其中一个tensor的维数等于1,要么两个维数相同。


CUDA语法

pytorch可以通过一行代码实现数据在GPU和CPU之间的切换,并且在网络训练过程中,你可以随时将数据从GPU中拿出来进行操作,包括打印之类的,比tf要方便一些。例如:

x = torch.cuda.FloatTensor(1)
# x.get_device() == 0
y = torch.FloatTensor(1).cuda()
# y.get_device() == 0

with torch.cuda.device(1):
    # allocates a tensor on GPU 1
    a = torch.cuda.FloatTensor(1)

    # transfers a tensor from CPU to GPU 1
    b = torch.FloatTensor(1).cuda()
    # a.get_device() == b.get_device() == 1

    c = a + b
    # c.get_device() == 1

    z = x + y
    # z.get_device() == 0

    # even within a context, you can give a GPU id to the .cuda call
    d = torch.randn(2).cuda(2)
    # d.get_device() == 2
pytoch可以允许用户指定特定的GPU来创建tensor,一旦创建就只能在当前设备下使用,跨GPU的操作目前是不允许的。

在多GPU环境中,充分利用GPU性能可以使用nn.DataParallel来将大的batch切成几份,分别送到不同GPU中的模型中进行训练,若你的batch只有1,却仍然要使用多GPU的话,目前来说是没有办法将一个模型运行在几个GPU上的。


序列化语法
pytorch中存储模型也非常简单。有两种方法:
1.只存储模型参数。这种方法得到的文件更小,再随后导入模型时需要使用与之前相同的模型参数来构建网络然后初始化参数。
torch.save(the_model.state_dict(), PATH)
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))
2.存储全部模型,包括计算图结构。
torch.save(the_model, PATH)
the_model = torch.load(PATH)

你可能感兴趣的:(Pytorch)