MMdet Model for Image Segmentation

6c9ac8f over 3 years ago

4.26 kB

import numpy as np
import pytest
import torch

from mmdet.utils import AvoidOOM
from mmdet.utils.memory import cast_tensor_type


def test_avoidoom():
    tensor = torch.from_numpy(np.random.random((20, 20)))
    if torch.cuda.is_available():
        tensor = tensor.cuda()
        # get default result
        default_result = torch.mm(tensor, tensor.transpose(1, 0))

        # when not occurred OOM error
        AvoidCudaOOM = AvoidOOM()
        result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
                                                          tensor.transpose(
                                                              1, 0))
        assert default_result.device == result.device and \
               default_result.dtype == result.dtype and \
               torch.equal(default_result, result)

        # calculate with fp16 and convert back to source type
        AvoidCudaOOM = AvoidOOM(test=True)
        result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
                                                          tensor.transpose(
                                                              1, 0))
        assert default_result.device == result.device and \
               default_result.dtype == result.dtype and \
               torch.allclose(default_result, result, 1e-3)

        # calculate on cpu and convert back to source device
        AvoidCudaOOM = AvoidOOM(test=True)
        result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
                                                          tensor.transpose(
                                                              1, 0))
        assert result.dtype == default_result.dtype and \
               result.device == default_result.device and \
               torch.allclose(default_result, result)

        # do not calculate on cpu and the outputs will be same as input
        AvoidCudaOOM = AvoidOOM(test=True, to_cpu=False)
        result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
                                                          tensor.transpose(
                                                              1, 0))
        assert result.dtype == default_result.dtype and \
               result.device == default_result.device

    else:
        default_result = torch.mm(tensor, tensor.transpose(1, 0))
        AvoidCudaOOM = AvoidOOM()
        result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
                                                          tensor.transpose(
                                                              1, 0))
        assert default_result.device == result.device and \
               default_result.dtype == result.dtype and \
               torch.equal(default_result, result)


def test_cast_tensor_type():
    inputs = torch.rand(10)
    if torch.cuda.is_available():
        inputs = inputs.cuda()
    with pytest.raises(AssertionError):
        cast_tensor_type(inputs, src_type=None, dst_type=None)
    # input is a float
    out = cast_tensor_type(10., dst_type=torch.half)
    assert out == 10. and isinstance(out, float)
    # convert Tensor to fp16 and re-convert to fp32
    fp16_out = cast_tensor_type(inputs, dst_type=torch.half)
    assert fp16_out.dtype == torch.half
    fp32_out = cast_tensor_type(fp16_out, dst_type=torch.float32)
    assert fp32_out.dtype == torch.float32

    # input is a list
    list_input = [inputs, inputs]
    list_outs = cast_tensor_type(list_input, dst_type=torch.half)
    assert len(list_outs) == len(list_input) and \
           isinstance(list_outs, list)
    for out in list_outs:
        assert out.dtype == torch.half
    # input is a dict
    dict_input = {'test1': inputs, 'test2': inputs}
    dict_outs = cast_tensor_type(dict_input, dst_type=torch.half)
    assert len(dict_outs) == len(dict_input) and \
           isinstance(dict_outs, dict)

    # convert the input tensor to CPU and re-convert to GPU
    if torch.cuda.is_available():
        cpu_device = torch.empty(0).device
        gpu_device = inputs.device
        cpu_out = cast_tensor_type(inputs, dst_type=cpu_device)
        assert cpu_out.device == cpu_device

        gpu_out = cast_tensor_type(inputs, dst_type=gpu_device)
        assert gpu_out.device == gpu_device

	import numpy as np
	import pytest
	import torch

	from mmdet.utils import AvoidOOM
	from mmdet.utils.memory import cast_tensor_type


	def test_avoidoom():
	tensor = torch.from_numpy(np.random.random((20, 20)))
	if torch.cuda.is_available():
	tensor = tensor.cuda()
	# get default result
	default_result = torch.mm(tensor, tensor.transpose(1, 0))

	# when not occurred OOM error
	AvoidCudaOOM = AvoidOOM()
	result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
	tensor.transpose(
	1, 0))
	assert default_result.device == result.device and \
	default_result.dtype == result.dtype and \
	torch.equal(default_result, result)

	# calculate with fp16 and convert back to source type
	AvoidCudaOOM = AvoidOOM(test=True)
	result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
	tensor.transpose(
	1, 0))
	assert default_result.device == result.device and \
	default_result.dtype == result.dtype and \
	torch.allclose(default_result, result, 1e-3)

	# calculate on cpu and convert back to source device
	AvoidCudaOOM = AvoidOOM(test=True)
	result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
	tensor.transpose(
	1, 0))
	assert result.dtype == default_result.dtype and \
	result.device == default_result.device and \
	torch.allclose(default_result, result)

	# do not calculate on cpu and the outputs will be same as input
	AvoidCudaOOM = AvoidOOM(test=True, to_cpu=False)
	result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
	tensor.transpose(
	1, 0))
	assert result.dtype == default_result.dtype and \
	result.device == default_result.device

	else:
	default_result = torch.mm(tensor, tensor.transpose(1, 0))
	AvoidCudaOOM = AvoidOOM()
	result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(tensor,
	tensor.transpose(
	1, 0))
	assert default_result.device == result.device and \
	default_result.dtype == result.dtype and \
	torch.equal(default_result, result)


	def test_cast_tensor_type():
	inputs = torch.rand(10)
	if torch.cuda.is_available():
	inputs = inputs.cuda()
	with pytest.raises(AssertionError):
	cast_tensor_type(inputs, src_type=None, dst_type=None)
	# input is a float
	out = cast_tensor_type(10., dst_type=torch.half)
	assert out == 10. and isinstance(out, float)
	# convert Tensor to fp16 and re-convert to fp32
	fp16_out = cast_tensor_type(inputs, dst_type=torch.half)
	assert fp16_out.dtype == torch.half
	fp32_out = cast_tensor_type(fp16_out, dst_type=torch.float32)
	assert fp32_out.dtype == torch.float32

	# input is a list
	list_input = [inputs, inputs]
	list_outs = cast_tensor_type(list_input, dst_type=torch.half)
	assert len(list_outs) == len(list_input) and \
	isinstance(list_outs, list)
	for out in list_outs:
	assert out.dtype == torch.half
	# input is a dict
	dict_input = {'test1': inputs, 'test2': inputs}
	dict_outs = cast_tensor_type(dict_input, dst_type=torch.half)
	assert len(dict_outs) == len(dict_input) and \
	isinstance(dict_outs, dict)

	# convert the input tensor to CPU and re-convert to GPU
	if torch.cuda.is_available():
	cpu_device = torch.empty(0).device
	gpu_device = inputs.device
	cpu_out = cast_tensor_type(inputs, dst_type=cpu_device)
	assert cpu_out.device == cpu_device

	gpu_out = cast_tensor_type(inputs, dst_type=gpu_device)
	assert gpu_out.device == gpu_device