I ran the code using dataparellel and got the error.
My demo code is as follows:
import math
import torch
import torch.nn as nn
class EnergyConstrainedADMM_P_imagenet(nn.Module):
def __init__(self, net_model, width_ub, blocks, scale=None):
super(EnergyConstrainedADMM_P_imagenet, self).__init__()
if scale != None:
width_ub = [width_ub[0]] + [int(i * scale) for i in width_ub[1:-1]] + [width_ub[-1]]
print(width_ub)
self.net = net_model
self.s = nn.Parameter(torch.FloatTensor(width_ub))
self.t1 = nn.Parameter(torch.rand(1, blocks[0]))
self.t2 = nn.Parameter(torch.rand(1, blocks[1]))
self.t3 = nn.Parameter(torch.rand(1, blocks[2]))
self.t4 = nn.Parameter(torch.rand(1, blocks[3]))
self.t = [self.t1, self.t2, self.t3, self.t4] #建立多个1维parameter
self.tau = 5.0
# 代码来源：TAS SoftSelect.py
def select2withP(self, logits, tau):
if tau <= 0:
probs = nn.functional.softmax(logits, dim=1)
else:
while True: # a trick to avoid the gumbels bug
gumbels = -torch.empty_like(logits).exponential_().log()
new_logits = (logits.log_softmax(dim=1) + gumbels) / tau
probs = nn.functional.softmax(new_logits, dim=1)
if (not torch.isinf(gumbels).any()) and (not torch.isinf(probs).any()) and (
not torch.isnan(probs).any()):
break
return probs
def set_tau(self, tau_max, tau_min, epoch_ratio):
self.tau = tau_min + (tau_max - tau_min) * (1 + math.cos(math.pi * epoch_ratio)) / 2
def forward(self, x):
probs = []
for i in range(len(self.t)):
probs.append(self.select2withP(self.t[i], self.tau))
features = []
for i, layer in enumerate(self.net.layers):
layer_i = layer(x)
features.append(layer_i)
# print('i:{}, layer_i:{}'.format(i, layer_i.shape))
if i in self.net.info:
choices = self.net.info[i]
tensors = features[choices[0]:choices[-2] + 1]
probs_i = probs[choices[-1]].view(-1)
x = sum([tensor * w for tensor, w in zip(tensors, probs_i)])
elif i == len(self.net.layers) - 2:
x = layer_i.view(layer_i.size(0), -1) # avgpool
else:
x = layer_i
return x
'''
bottleneck block, which is used for ResNet-50,101
'''
class Bottleneck(nn.Module):
def __init__(self, inplanes, cfg, stride=1):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, cfg[0], kernel_size=1, bias=False)
self.conv2 = nn.Conv2d(cfg[0], cfg[1], kernel_size=3, stride=stride, padding=1, bias=False)
self.conv3 = nn.Conv2d(cfg[1], cfg[2], kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(cfg[0])
self.bn2 = nn.BatchNorm2d(cfg[1])
self.bn3 = nn.BatchNorm2d(cfg[2])
self.relu = nn.ReLU(inplace=True)
self.downsample = None
if stride != 1 or inplanes != cfg[2]:
self.downsample = nn.Sequential(
nn.Conv2d(inplanes, cfg[2], kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(cfg[2]),
)
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class ResNet_Bottleneck(nn.Module):
def __init__(self, cfg, nums, Class_num=1000):
self.cfg = cfg
self.inplanes = 64
super(ResNet_Bottleneck, self).__init__()
self.layers = nn.ModuleList([nn.Sequential(
nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)])
strides = [1, 2, 2, 2]
self.info = {}
for i in range(4): # 4 stage
filters = cfg[sum(nums[:i]) * 3: sum(nums[:i + 1]) * 3]
self.layers.append(Bottleneck(self.inplanes, filters[0:3], strides[i]))
self.inplanes = filters[2]
for num in range(1, nums[i]):
self.layers.append(Bottleneck(self.inplanes, filters[num * 3:(num + 1) * 3]))
self.inplanes = filters[num * 3 + 2]
self.info[sum(nums[:i + 1])] = list(range(sum(nums[:i]) + 1, sum(nums[:i + 1]) + 1)) + [i]
self.layers.append(nn.AvgPool2d(7, stride=1))
self.layers.append(nn.Linear(cfg[-1], Class_num))
# print(self.info)
# 初始化
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.bias.data.zero_()
def forward(self, x):
for i, layer in enumerate(self.layers):
layer_i = layer(x)
if i == len(self.layers) - 2:
x = layer_i.view(layer_i.size(0), -1) # avgpool
else:
x = layer_i
return x
def get_bn_weights(self):
res = []
for m in self.modules():
if isinstance(m, Bottleneck):
res += [m.bn1.weight, m.bn2.weight, m.bn3.weight]
return res
def bn_update(self, args):
for m in self.modules():
if isinstance(m, Bottleneck):
m.bn1.weight.grad.data.add_(args.bn_sparsity * torch.sign(m.bn1.weight.data)) # L1
m.bn2.weight.grad.data.add_(args.bn_sparsity * torch.sign(m.bn2.weight.data))
m.bn3.weight.grad.data.add_(args.bn_sparsity * torch.sign(m.bn3.weight.data))
def resnet50(cfg=None):
if cfg is None:
cfg = [[64, 64, 256] * 3, [128, 128, 512] * 4, [256, 256, 1024] * 6, [512, 512, 2048] * 3]
cfg = [item for sub_list in cfg for item in sub_list]
nums = [3, 4, 6, 3]
model = ResNet_Bottleneck(cfg, nums)
return model
#get model
model = resnet50().cuda()
width_ub = [3, 64] + resnet50().cfg + [1000]
blocks = [3, 4, 6, 3]
primal_model = EnergyConstrainedADMM_P_imagenet(model, width_ub, blocks).cuda()
primal_model = torch.nn.DataParallel(primal_model)
#test
img = torch.rand(4,3,224,224).cuda()
out = model(img)
out = primal_model(img)
target = torch.tensor([1,1,1,1]).cuda()
criterion = nn.CrossEntropyLoss()
loss = criterion(out, target)
loss.backward()
the error is as follows:
Traceback (most recent call last):
File "demo.py", line 210, in <module>
loss.backward()
File "/home/lth/anaconda3/lib/python3.7/site-packages/torch/tensor.py", line 118, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/home/lth/anaconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 93, in backward
allow_unreachable=True) # allow_unreachable flag
RuntimeError: Function MulBackward0 returned an invalid gradient at index 1 - expected device cuda:0 but got cuda:2
