Five PyTorch functions recommended for machine learning beginners

PyTorch is a free and opensource machine learning library based on the Torch library. It has numerous in-built functions that allows data scientists to build machine learning algorithms with the shortest lines of code. In this article, I will outline and explain five interesting functions that every beginner in the field of data science will find useful when it comes to building models using PyTorch. These functions include;

  • torch.randn()
  • torch.split()
  • torch.reshape()
  • torch.transpose()

Before we begin, let’s install and import PyTorch

Function 1 —

The function concatenates two or more tensors of the same dimension into a single result tensor.

This example shows the use of to join two tensors into a single output tensor.

In this second example, we use function with an optional dim=dimension argument. Expected dimension values are 0 and 1. The 0 dimension makes the concatenation to occur along the vertical axis where as the 1 dimension makes the concatenation to occur along the horizontal axis. If not specified, the default dimension is 0 and all specified tensor are concatenated along the vertical axis.

In example 3, the function fails to execute simply because of invalid argument combinations provided. The function expect a tuple of tensors and not individual tensors as arguments. It is easy to forget to pass the tensors as a tuple hence failure to execute.

For beginners, joining of tensors will be a key requirement in manipulation of datasets hence knowing how to and when to do the same will be an important step in the right direction.

Function 2 — torch.randn()

THe torch.randn() is used to generate a tensor filled with random numbers from a standard normal distribution. The function syntax is shown below:

Syntax: torch.randn(*size, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)

where the size argument is required argument which specifies the sequence of integers defining the size of the output tensor.

The rest of the function arguments are optional.

The examples below demonstrate the use of torch.randn().

Example 1 above shows the basic use of torch.randn() to generate a tensor of the shape 2,3,2. The values are randomly generated and contains both positive and negative values.

The torch.randn() function in example two returns a 4 x 3 x 2 tensor where also the gradient is recorded for further use.

Example 3 demonstrates that when using torch.randn(), at least one argument must be provided for the function to correctly execute and return results. This is important because the returned tensor will need to have an expected shape for it to be useful and that shape must be specified in advance.

The torch.randn() function is a very useful function that every beginner must learn. It is mostly applied in generation of initial weights matrices and bias vectors that are necessary for training of model across different machine learning algorithms.

Function 3 — torch.split()

The torch.split() Splits the tensor into chunks based on the specified split size. The function accepts the arguments as specified in the syntax below.

Syntax: torch.split(tensor, split_size_or_sections, dim=0)

where tensor is the input tensor to be splitted, split_size_or_sections is either an integer number or a list specified the split chunk size and the dim argument specifies the dimension on which to split and this can be either 0 or 1.

In this example, we have generated an 8 x 2 tensor and managed to split the tensor into two chunks using the size argument value 4.

It is important to note that when specifying the size argument, the value specified refers to the size of the output tensor chunk and not the number of chunks returned.

In example 2 above, we are splitting a tensor by specifying a lift of sections to split at. It is important to note that the sum of specified values in the list must be equal to the length of the input tensor corresponding to the split dimension.

For example, the shape of the input tensor is 4 x 2 and therefore the contents of the split list [1,3] must add up to 4. The first output will have one row and two columns where as the last output will have three rows and two columns.

In the example 3, we have an input tensor with a 4 x 2 shape and we have specified [2,1] as the split sections list. We are getting an error because our the sum of values in our splitting list does not equal the length of the list corresponding to the split dimension which is 1.

If we specify [3,1] as the splitting list, we would expect our function to execute given that 3+1=4 will be equal to 4 which is the length of the input list which we intend to split along the dim=1.

The torch.split() is an important function to understand as it is quite useful in generation, testing and training of machine learning models. In most cases, an input data will need to be split into training and testing datasets and most commonly, the torch.split() will be the function to use.

Function 4 — torch.reshape()

The torch.reshape() function returns a tensor with the same data and number of elements as input, but with the specified shape. The function simply changes the shape of a tensor from the original shape to the specified shape.

In this example, we generate a 1 x 10 tensor using the torch.arange() function and then we reshape the tensor into a 2 x 5 tensor. It is important to note that the specified shape must be compatible with the original shape for the function to return the desired result without generating an error.

In example 2 above, we generate a 3 dimension tensor with the shape (3,1,4) and reshape it into a (3,2,2) tensor. The most important factor here is to ensure that the shape of the resultant tensor is compatible with the input tensor.

In example 3, we try to reshape a (1,10) tensor into a (4,2) tensor and we get an error message because the two provided shapes are not compatible.

Reshaping a tensor from one shape to another is a useful operation for machine learning tasks. The point to note is the fact that the input and output tensor shapes must be compatible in other words, the product of the elements in the provided shapes for both input and output tensors must always be the same.

Function 5 — torch.transpose()

The torch.transpose() is used to transpose the values of a tensor from one specified dimension to another. This operation is useful in the multiplication of matrices and vectors. The syntax for using the transpose function is shown below:

syntax: torch.transpose(input, dim0, dim1)

where input is the original tensor to be transposed, dim0 is the dimension to transpose from, and dim1 is the dimension to transpose to.

In this simple example, we transpose a 2 dimensional 2 x 3 tensor generated by the randn() function from the horizontal axis into vertical axis i.e from dim=0 to dim=1. The result is such that the contents of the first row, becomes the contents in the first column in the resulting tensor.

In example 2 above, we are transposing a 3 dimensional matrix with the shape 2x3x2 from horizontal to vertical axis. An interesting comparison of the input and output tensors shows that when transposing a 3 dimensional matrix, the matrix is reshaped such that if the shape of the input tensor is (x,y,z) then the resulting tensor will have a shape with dimensions (y,x,z).

The difference between transpose and reshape in this case is the arrangement of values in the resulting tensor.

Example three demonstrates that all the arguments for torch.transpose() function are required and must be specified for successful execution of a pytorch code containing a torch.transpose() function.

As mentioned already, torch.transpose() is one of the key pytorch functions that every data science beginner ought to understand due to its wide application in machine learning model development.


In this notebook, the focus was on the five most important pytorch functions that every data science beginner ought to learn. Starting with the, we have demonstrated that the importance of joining tensors either on the vertical or horizontal axis.

We have also extensively used the pytorch.randn() function through out this notebook for generating sample tensors in various examples. This shows that, a beginner will not have to manually type sample tensors for learning purposes but will quickly be able generated sample to quicken the learning process.

We have also seen how to split tensors which is the opposite of concatenation of tensors as well as reshaping and transposing tensors. I do hope that by reading through this article, one will be able to quickly grasp the concepts illustrated by the numerous examples here in.

Reference Links

Provide links to your references and other interesting articles about tensors

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store