A slide of memory efficient pytorch including inplace, memory sharing and re-computation tricks.

1. Memory Efficient Pytorch
SNU RPLab
Hyungjoo Cho

2. Computation Graph
w₁ x₁ w₂ x₂ b
z
h
L
y

3. Computation Graph
w₁ x₁ w₂ x₂ b
z
h
L
y

4. Computation Graph
w₁ x₁ w₂ x₂ b
z
h
yL

5. Computation Graph
w₁ x₁ w₂ x₂ b
z
h
yL

6. Computation Graph
w₁ x₁ w₂ x₂
z
h
yL
b

7. Network Configuration
INPUT
FC
Sigmoid
LogSigmoid Label

8. Gradient Calculation Graph
INPUT
FC-forward
Sigm-forward
LogSigm-forward LogSigm-backward
FC-backward
Sigm-backward
Label
INPUT-Grad

9. Memory Allocation
INPUT
FC-forward
Sigm-forward
LogSigm-forward LogSigm-backward
FC-backward
Sigm-backward
Label
INPUT-Grad

10. Efficient Memory Allocation
INPUT
FC-forward
Sigm-forward
LogSigm-forward LogSigm-backward
FC-backward
Sigm-backward
Label
INPUT-Grad

11. Efficient Memory Allocation
INPUT
FC-forward
Sigm-forward
LogSigm-forward LogSigm-backward
FC-backward
Sigm-backward
Label
INPUT-Grad
①
②
1. refcount(①) is 1.
2. size(output(①))
is same as
size(output(②)).

12. Efficient Memory Allocation
INPUT
FC-forward
Sigm-forward
LogSigm-forward LogSigm-backward
FC-backward
Sigm-backward
Label
INPUT-Grad
①
②
1. refcount(①) is 1.
2. size(output(①))
is same as
size(output(②)).
In-place operation

13. In-place Operation

14. Sorting
Bubble Merge

15. In-place Operation
• The input is usually overwritten by the output as the algorithm executes.
• In-place operation updates input sequence only through replacement or
swapping of elements.

16. In-place in Pytorch

17. a = torch.zeros(1, 3)
In-place in Pytorch

18. a = torch.zeros(1, 3)
print(a)
print(hex(id(a)))
In-place in Pytorch

19. a = torch.zeros(1, 3)
print(a)
print(hex(id(a)))
0 0 0
[torch.FloatTensor of size 1x3]
0x7f4b08813188
In-place in Pytorch

20. a = torch.zeros(1, 3)
print(a)
print(hex(id(a)))
0 0 0
[torch.FloatTensor of size 1x3]
0x7f4b08813188
0x7f4b08813188 FloatTensor([0, 0, 0])a
In-place in Pytorch

21. case 1)
a = a + 1
print(a.numpy())
print(hex(id(a)))
0x7f4b08813188 FloatTensor([0, 0, 0])a
case 2)
for i in range(3):
a[:, i] += 1
print(a.numpy())
print(hex(id(a)))
In-place in Pytorch

22. case 1)
a = a + 1
print(a.numpy())
print(hex(id(a)))
[[1. 1. 1.]]
0x7f4b088135c8
0x7f4b08813188 FloatTensor([0, 0, 0])a
case 2)
for i in range(3):
a[:, i] += 1
print(a.numpy())
print(hex(id(a)))
[[1. 1. 1.]]
0x7f4b08813188
In-place in Pytorch

23. case 1)
0x7f4b08813188 FloatTensor([0, 0, 0])a
case 2)
0x7f4b08813188 FloatTensor([0, 0, 0])
a
0x7f4b088135c8 FloatTensor([1, 1, 1])
0x7f4b08813188
FloatTensor([0, 0, 0])
FloatTensor([1, 1, 1])
a
In-place in Pytorch

24. case 1)
0x7f4b08813188 FloatTensor([0, 0, 0])a
case 2)
0x7f4b08813188 FloatTensor([0, 0, 0])
a
0x7f4b088135c8 FloatTensor([1, 1, 1])
0x7f4b08813188
FloatTensor([0, 0, 0])
FloatTensor([1, 1, 1])
a
Out of place In-place
In-place in Pytorch

25. case 3)
a = a.add(1)
print(a.numpy())
print(hex(id(a)))
0x7f4b08813188 FloatTensor([0, 0, 0])a
case 4)
a.add_(1)
print(a.numpy())
print(hex(id(a)))
In-place in Pytorch

26. case 3)
a = a.add(1)
print(a.numpy())
print(hex(id(a)))
[[1. 1. 1.]]
0x7f4b088135c8
0x7f4b08813188 FloatTensor([0, 0, 0])a
case 4)
a.add_(1)
print(a.numpy())
print(hex(id(a)))
[[1. 1. 1.]]
0x7f4b08813188
In-place in Pytorch
Out of place In-place

27. case 5)
a += 1
print(a.numpy())
print(hex(id(a)))
0x7f4b08813188 FloatTensor([0, 0, 0])a
In-place in Pytorch

28. case 5)
a += 1
print(a.numpy())
print(hex(id(a)))
[[1. 1. 1.]]
0x7f4b08813188
0x7f4b08813188 FloatTensor([0, 0, 0])a
In-place in Pytorch
In-place

29. case 5)
a += 1
print(a.numpy())
print(hex(id(a)))
[[1. 1. 1.]]
0x7f4b08813188
0x7f4b08813188 FloatTensor([0, 0, 0])a
In-place in Pytorch
In-place
torch/autograd/variable.py::Variable( )

30. a = a.fill(3)
In-place in Pytorch

31. a = a.fill(3)
AttributeError: ‘torch.FloatTensor’
object has no attribute ‘fill’
In-place in Pytorch

32. a.fill_(3)
print(a.numpy())
[[3. 3. 3.]]
In-place in Pytorch

33. In-place in Pytorch
In-place ‘only’
- fill_()
- zero_()
- normal_()
- uniform_()
- exponential_()
- etc…

34. In-place in Pytorch
a = a.t_()
3.
3.
3.
[torch.FloatTensor of size 3x1]
Different tensor size, but same buffer size

35. Non-linear Activation
self.add_module(‘conv’, conv2d_3x3(in_dim, out_dim))
self.add_module(‘bn’, nn.BatchNorm2d(in_dim)
self.add_module(‘act’, nn.ReLU(inplace=True))

36. Non-linear Activation
void THNN_(Threshold_updateOutput)(THNNState *state,
THTensor *input,
THTensor *output,
accreal threshold_,
accreal val_,
bool inplace)
{
float threshold = (float)threshold_;
float val = (float)val_;
if (inplace)
{
int TH_TENSOR_APPLY_hasFinished = 0;
int64_t TH_TENSOR_dim_index = 0;
TH_TENSOR_APPLYX_PREAMBLE(float, input, -1, 0);
while (!TH_TENSORAPPLY_hasFinished)
{
for (; input_i < input_size; input_i++, input_data += input_stride)
{
if(*input_data <= threshold)
*input_data = val;
}
__TH_TENSOR_APPLYX_UPDATE_COUNTERS(TENSOR, 1);
}
THFree(input_counter);
THFloatTensor_set(output, input);
}
else
{
THFloatTensor_resizeAs(output, input);
int TH_TENSOR_APPLY_hasFinished = 0;
int64_t TH_TENSOR_dim_index = 0;
TH_TENSOR_APPLYX_PREAMBLE(float, output, -1, 0);
TH_TENSOR_APPLYX_PREAMBLE(float, input, -1, 0);
if(output_n != input_n)
{
THDescBuff T1buff = _THSizeDesc(output->size, output->nDimension);
THDescBuff T2buff = _THSizeDesc(input->size, input->nDimension);
}
while (!TH_TENSOR_APPLY_hasFinished)
{
for (; output_i < output_size && input_i < input_size; output_i++,
input_++, output_data += output_stride, input_data += input_stride)
{
*output_data = (*input_data > threshold) ? *input_data : val;
}
__TH_TENSOR_APPLYX_UPDATE_COUNTERS(output, 0);
__TH_TENSOR_APPLYX_UPDATE_COUNTERS(input, 0);
}
if(output_counter != NULL)
THFree(output_counter);
if(input_counter != NULL)
THFree(input_counter);
}
}

37. Non-linear Activation
void THNN_(Threshold_updateOutput)(THNNState *state,
THTensor *input,
THTensor *output,
accreal threshold_,
accreal val_,
bool inplace)
{
float threshold = (float)threshold_;
float val = (float)val_;
if (inplace)
{
int TH_TENSOR_APPLY_hasFinished = 0;
int64_t TH_TENSOR_dim_index = 0;
TH_TENSOR_APPLYX_PREAMBLE(float, input, -1, 0);
while (!TH_TENSORAPPLY_hasFinished)
{
for (; input_i < input_size; input_i++, input_data += input_stride)
{
if(*input_data <= threshold)
*input_data = val;
}
__TH_TENSOR_APPLYX_UPDATE_COUNTERS(TENSOR, 1);
}
THFree(input_counter);
THFloatTensor_set(output, input);
}
else
{
THFloatTensor_resizeAs(output, input);
int TH_TENSOR_APPLY_hasFinished = 0;
int64_t TH_TENSOR_dim_index = 0;
TH_TENSOR_APPLYX_PREAMBLE(float, output, -1, 0);
TH_TENSOR_APPLYX_PREAMBLE(float, input, -1, 0);
if(output_n != input_n)
{
THDescBuff T1buff = _THSizeDesc(output->size, output->nDimension);
THDescBuff T2buff = _THSizeDesc(input->size, input->nDimension);
}
while (!TH_TENSOR_APPLY_hasFinished)
{
for (; output_i < output_size && input_i < input_size; output_i++,
input_++, output_data += output_stride, input_data += input_stride)
{
*output_data = (*input_data > threshold) ? *input_data : val;
}
__TH_TENSOR_APPLYX_UPDATE_COUNTERS(output, 0);
__TH_TENSOR_APPLYX_UPDATE_COUNTERS(input, 0);
}
if(output_counter != NULL)
THFree(output_counter);
if(input_counter != NULL)
THFree(input_counter);
}
}
In-place Out of place

38. Non-linear Activation
A
C
B
INPUT
Sigmoid(A)
Sigmoid(B)
D FPool(C) D + E
E
B
Pool(B)

39. Reference Count
A
C
B
INPUT
Sigmoid(A)
Sigmoid(B)
D FPool(C) D + E
E
B
Pool(B)
Node refcount
A 1
B 2
C 1
D 1
E 1
F 1

40. Mark Dirty
A
C
B
INPUT
Sigmoid(A)
Sigmoid(B)
D FPool(C) D + E
E
B
Pool(B)
Node refcount
A 1
B 2
C 1
D 1
E 1
F 1
If B is In-place operator,
mark_dirty( ) raises an error.

41. Memory sharing
A
B
INPUT
Sigmoid(A)
Sigmoid(B)
FPool(C) D + E
B
Pool(B)C
D
E

42. Memory sharing
A
B
INPUT
Sigmoid(A)
Sigmoid(B)
FPool(C) D + E
B
Pool(B)C
D
ERe-use
Release B after allocating C, E
Reuse for D
Memory sharing : Memory used by intermediate results that are no longer needed can be recycled and used in another node.

43. Memory sharing
A
B
INPUT
Sigmoid(A)
Sigmoid(B)
FPool(C) D + E
B
Pool(B)C
D
E
In-place
A
B
INPUT
Sigmoid(A)
Sigmoid(B)
FPool(C) D + E
B
Pool(B)C
D
E
In-place
OR A
B
INPUT
Sigmoid(A)
Sigmoid(B)
FPool(C) D + E
B
Pool(B)C
D
E
OR
Re-use

44. Trade Computation for Memory
• Apply normalization and non-linearities before/after the conv-operation.
• Convolution is most efficient when input lies in a contiguous block of
memory
• To make a contiguous input, each layer must copy all previous features
(concatenation → mem-copy)
• Above operations are computationally extremely cheap
• Copying to pre-allocated memory is significantly faster than
allocating new memory

45. Shared storage for concatenation
• Rather than allocating memory for each concatenation operation, assign
the outputs to a memory allocation shared across all layers
• Shared memory storage is used by all network layers, its data is not
permanent
• Need to be recomputed during back-propagation

46. Shared storage for batch normalization
& non-linearity activation
• Assign the outputs of batch normalization / activation to a shared
memory allocation
• The data in shared memory storage is not permanent and will be
overwritten by the next layer
• Should recompute the batch normalization / activation outputs during
back-propagation

47. Re-computation
INPUT
conv-forward
bn-forward
relu-forward
conv-forward
bn-forward
relu-forward

48. Re-computation
INPUT
conv-forward
bn-forward
relu-forward
conv-forward
bn-forward
relu-forward
conv-backward
bn-backward
relu-backward
conv-backward
bn-backward
relu-backward
INPUT-Grad

49. Re-computation
INPUT
conv-forward
bn-forward
relu-forward
conv-forward
bn-forward
relu-forward
conv-backward
bn-backward
relu-backward
conv-backward
bn-backward
relu-backward
INPUT-Grad

50. Result
ResNet

51. Result
DenseNet

52. Benefit
• Can increase mini-batch size
→ Speed up
• Build deeper model
→ Accuracy up
• Can use deep model using small GPU
→ Money up

53. Thanks 😊