Python源码示例：torch.Storage()
示例1
def __init__(self, in_num, out_num, layer_num, max_link, storage_size=1024):
input_storage_1 = torch.Storage(storage_size)
input_storage_2 = torch.Storage(storage_size)
self.shared_allocation_1 = _SharedAllocation(input_storage_1)
self.shared_allocation_2 = _SharedAllocation(input_storage_2)
max_in_num = in_num + out_num * max_link
self.final_num_features = max_in_num
self.saved_features = []
self.max_link = max_link
super(_IntermediaBlock, self).__init__()
print('creating intermedia block ...')
self.adapters = []
for i in range(0, layer_num-1):
if i < max_link:
tmp_in_num = in_num + (i+1) * out_num
else:
tmp_in_num = max_in_num
print('intermedia layer %d input channel number is %d' % (i, tmp_in_num))
self.adapters.append(_EfficientDensenetBottleneck(self.shared_allocation_1,
self.shared_allocation_2,
tmp_in_num, out_num))
self.adapters = nn.ModuleList(self.adapters)
print('intermedia layer output channel number is %d' % out_num)
示例2
def __init__(self, in_num, out_num, layer_num, max_link, storage_size=1024):
input_storage_1 = torch.Storage(storage_size)
input_storage_2 = torch.Storage(storage_size)
self.shared_allocation_1 = _SharedAllocation(input_storage_1)
self.shared_allocation_2 = _SharedAllocation(input_storage_2)
max_in_num = in_num + out_num * max_link
self.final_num_features = max_in_num
self.saved_features = []
self.max_link = max_link
super(_IntermediaBlock, self).__init__()
print('creating intermedia block ...')
self.adapters = []
for i in range(0, layer_num-1):
if i < max_link:
tmp_in_num = in_num + (i+1) * out_num
else:
tmp_in_num = max_in_num
print('intermedia layer %d input channel number is %d' % (i, tmp_in_num))
self.adapters.append(_EfficientDensenetBottleneck(self.shared_allocation_1,
self.shared_allocation_2,
tmp_in_num, out_num))
self.adapters = nn.ModuleList(self.adapters)
print('intermedia layer output channel number is %d' % out_num)
示例3
def _map_location(xpu, storage, location):
"""
Helper for `xpu.load` used when calling `torch.load`
Args:
storage (torch.Storage) : the initial deserialization of the
storage of the data read by `torch.load`, residing on the CPU.
location (str): tag identifiying the location the data being read
by `torch.load` was originally saved from.
Returns:
torch.Storage : the storage
"""
if xpu.is_gpu():
return storage.cuda(xpu._main_device_id)
else:
return storage
示例4
def device_mapping(cuda_device: int):
"""
In order to `torch.load()` a GPU-trained model onto a CPU (or specific GPU),
you have to supply a `map_location` function. Call this with
the desired `cuda_device` to get the function that `torch.load()` needs.
"""
def inner_device_mapping(storage: torch.Storage, location) -> torch.Storage:
if cuda_device >= 0:
return storage.cuda(cuda_device)
else:
return storage
return inner_device_mapping
示例5
def device_mapping(cuda_device: int):
"""
In order to `torch.load()` a GPU-trained model onto a CPU (or specific GPU),
you have to supply a `map_location` function. Call this with
the desired `cuda_device` to get the function that `torch.load()` needs.
"""
def inner_device_mapping(storage: torch.Storage, location) -> torch.Storage: # pylint: disable=unused-argument
if cuda_device >= 0:
return storage.cuda(cuda_device)
else:
return storage
return inner_device_mapping
示例6
def device_mapping(cuda_device: int):
"""
In order to `torch.load()` a GPU-trained model onto a CPU (or specific GPU),
you have to supply a `map_location` function. Call this with
the desired `cuda_device` to get the function that `torch.load()` needs.
"""
def inner_device_mapping(storage: torch.Storage, location) -> torch.Storage: # pylint: disable=unused-argument
if cuda_device >= 0:
return storage.cuda(cuda_device)
else:
return storage
return inner_device_mapping
示例7
def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, storage_size=1024):
input_storage_1 = torch.Storage(storage_size)
input_storage_2 = torch.Storage(storage_size)
self.final_num_features = num_input_features + (growth_rate * num_layers)
self.shared_allocation_1 = _SharedAllocation(input_storage_1)
self.shared_allocation_2 = _SharedAllocation(input_storage_2)
super(_DenseBlock, self).__init__()
for i in range(num_layers):
layer = _DenseLayer(self.shared_allocation_1, self.shared_allocation_2, num_input_features + i * growth_rate,
growth_rate, bn_size, drop_rate)
self.add_module('denselayer%d' % (i + 1), layer)
示例8
def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, storage_size=1024):
input_storage_1 = torch.Storage(storage_size)
input_storage_2 = torch.Storage(storage_size)
self.final_num_features = num_input_features + (growth_rate * num_layers)
self.shared_allocation_1 = _SharedAllocation(input_storage_1)
self.shared_allocation_2 = _SharedAllocation(input_storage_2)
super(_DenseBlock, self).__init__()
for i in range(num_layers):
layer = _DenseLayer(self.shared_allocation_1, self.shared_allocation_2, num_input_features + i * growth_rate,
growth_rate, bn_size, drop_rate)
self.add_module('denselayer%d' % (i + 1), layer)
示例9
def __init__(self, size):
self._cpu_storage = torch.Storage(size)
self._gpu_storages = []
if torch.cuda.is_available():
for device_idx in range(torch.cuda.device_count()):
with torch.cuda.device(device_idx):
self._gpu_storages.append(torch.Storage(size).cuda())
示例10
def device_mapping(cuda_device: int):
"""
In order to `torch.load()` a GPU-trained model onto a CPU (or specific GPU),
you have to supply a `map_location` function. Call this with
the desired `cuda_device` to get the function that `torch.load()` needs.
"""
def inner_device_mapping(storage: torch.Storage, location) -> torch.Storage: # pylint: disable=unused-argument
if cuda_device >= 0:
return storage.cuda(cuda_device)
else:
return storage
return inner_device_mapping
示例11
def device_mapping(cuda_device: int):
"""
In order to `torch.load()` a GPU-trained model onto a CPU (or specific GPU),
you have to supply a `map_location` function. Call this with
the desired `cuda_device` to get the function that `torch.load()` needs.
"""
def inner_device_mapping(storage: torch.Storage, location) -> torch.Storage: # pylint: disable=unused-argument
if cuda_device >= 0:
return storage.cuda(cuda_device)
else:
return storage
return inner_device_mapping
示例12
def __init__(self, in_num, neck_size, growth_rate, layer_num, max_link,
storage_size=1024, requires_skip=True, is_up=False):
input_storage_1 = torch.Storage(storage_size)
input_storage_2 = torch.Storage(storage_size)
self.shared_allocation_1 = _SharedAllocation(input_storage_1)
self.shared_allocation_2 = _SharedAllocation(input_storage_2)
self.saved_features = []
self.max_link = max_link
self.requires_skip = requires_skip
super(_DenseBlock, self).__init__()
max_in_num = in_num + max_link * growth_rate
self.final_num_features = max_in_num
self.layers = []
print('layer number is %d' % layer_num)
for i in range(0, layer_num):
if i < max_link:
tmp_in_num = in_num + i * growth_rate
else:
tmp_in_num = max_in_num
print('layer %d input channel number is %d' % (i, tmp_in_num))
self.layers.append(_DenseLayer(self.shared_allocation_1, self.shared_allocation_2,
tmp_in_num, neck_size, growth_rate))
self.layers = nn.ModuleList(self.layers)
self.adapters_ahead = []
adapter_in_nums = []
adapter_out_num = in_num
if is_up:
adapter_out_num = adapter_out_num / 2
for i in range(0, layer_num):
if i < max_link:
tmp_in_num = in_num + (i+1) * growth_rate
else:
tmp_in_num = max_in_num + growth_rate
adapter_in_nums.append(tmp_in_num)
print('adapter %d input channel number is %d' % (i, adapter_in_nums[i]))
self.adapters_ahead.append(_EfficientDensenetBottleneck(self.shared_allocation_1,
self.shared_allocation_2,
adapter_in_nums[i], adapter_out_num))
self.adapters_ahead = nn.ModuleList(self.adapters_ahead)
print('adapter output channel number is %d' % adapter_out_num)
if requires_skip:
print('creating skip layers ...')
self.adapters_skip = []
for i in range(0, layer_num):
self.adapters_skip.append(_EfficientDensenetBottleneck(self.shared_allocation_1,
self.shared_allocation_2,
adapter_in_nums[i], adapter_out_num))
self.adapters_skip = nn.ModuleList(self.adapters_skip)
示例13
def __init__(self, in_num, neck_size, growth_rate, layer_num, max_link,
storage_size=1024, requires_skip=True, is_up=False):
input_storage_1 = torch.Storage(storage_size)
input_storage_2 = torch.Storage(storage_size)
self.shared_allocation_1 = _SharedAllocation(input_storage_1)
self.shared_allocation_2 = _SharedAllocation(input_storage_2)
self.saved_features = []
self.max_link = max_link
self.requires_skip = requires_skip
super(_DenseBlock, self).__init__()
max_in_num = in_num + max_link * growth_rate
self.final_num_features = max_in_num
self.layers = []
print('layer number is %d' % layer_num)
for i in range(0, layer_num):
if i < max_link:
tmp_in_num = in_num + i * growth_rate
else:
tmp_in_num = max_in_num
print('layer %d input channel number is %d' % (i, tmp_in_num))
self.layers.append(_DenseLayer(self.shared_allocation_1, self.shared_allocation_2,
tmp_in_num, neck_size, growth_rate))
self.layers = nn.ModuleList(self.layers)
self.adapters_ahead = []
adapter_in_nums = []
adapter_out_num = in_num
if is_up:
adapter_out_num = adapter_out_num / 2
for i in range(0, layer_num):
if i < max_link:
tmp_in_num = in_num + (i+1) * growth_rate
else:
tmp_in_num = max_in_num + growth_rate
adapter_in_nums.append(tmp_in_num)
print('adapter %d input channel number is %d' % (i, adapter_in_nums[i]))
self.adapters_ahead.append(_EfficientDensenetBottleneck(self.shared_allocation_1,
self.shared_allocation_2,
adapter_in_nums[i], adapter_out_num))
self.adapters_ahead = nn.ModuleList(self.adapters_ahead)
print('adapter output channel number is %d' % adapter_out_num)
if requires_skip:
print('creating skip layers ...')
self.adapters_skip = []
for i in range(0, layer_num):
self.adapters_skip.append(_EfficientDensenetBottleneck(self.shared_allocation_1,
self.shared_allocation_2,
adapter_in_nums[i], adapter_out_num))
self.adapters_skip = nn.ModuleList(self.adapters_skip)
