def getblock_max(cnt, name):
if cnt == 0:
return 0
model = ""
block_max = 0
blockname = "BasicBlock"
sys.path.append("input")
if "resnet18" in name:
from ResNet import resnet18
model = resnet18()
elif "resnet34" in name:
from ResNet import resnet34
model = resnet34()
elif "resnet50" in name:
from ResNet import resnet50
model = resnet50()
blockname = "Bottleneck"
modeltxt = f'.debug/{name}.txt'
with open(modeltxt, 'w') as fmd:
fmd.write(str(model))
with open(modeltxt, 'r') as frd:
for line in frd:
line = line.strip()
if line.startswith(f'(layer{cnt}): Sequential('):
for line in frd:
line = line.strip()
if f"): {blockname}(" in line:
block_max = line.split(')')[0].split('(')[1]
elif line.startswith(f'(layer{cnt+1}): Sequential(') or "quant_fc" in line:
return int(block_max)+1
assert len(sys.argv) == 5, f'Usage: {sys.argv[0]} name/to/net path/to/image path/to/ptfile path/to/output/dir'
net_name = sys.argv[1]
img_path = sys.argv[2]
pt_path = sys.argv[3]
output_dir = sys.argv[4]
model = vgg()
if sys.argv[1] == "vggnet":
model = vgg(num_classes=1000)
elif sys.argv[1] == "resnet18":
from ResNet import resnet18
model = resnet18()
elif sys.argv[1] == "resnet34":
from ResNet import resnet34
model = resnet34()
elif sys.argv[1] == "resnet50":
from ResNet import resnet50
model = resnet50()
# TODO: merge this into vgg.py
for m in model.modules():
if isinstance(m, (torch.nn.Conv2d, torch.nn.Linear)):
m.register_buffer(f'scale', torch.tensor(0.0))
# if isinstance(m, torch.nn.BatchNorm2d):
# m.register_buffer(f'bn_k', torch.zeros_like(m.weight.data))
# m.register_buffer(f'bn_b', torch.zeros_like(m.weight.data))
# Load checkpoint
prefix = 'model.'
state_dict = torch.load(pt_path, map_location=torch.device('cpu'))['state_dict']
'''
#查看图片样本
test_data_iter = iter(validate_loader)
test_image, test_label = test_data_iter.next()
def imshow(img):
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
print(' '.join('%5s' % cla_dict[test_label[j].item()] for j in range(4)))
imshow(utils.make_grid(test_image))
'''
#-----------------------网络构建----------------------------
net = resnet34()
# load pretrain weights
model_weight_path = path + "./resnet34-pre.pth"
missing_keys, unexpected_keys = net.load_state_dict(
torch.load(model_weight_path), strict=False)
# for param in net.parameters():
# param.requires_grad = False
# change fc layer structure
inchannel = net.fc.in_features
net.fc = nn.Linear(inchannel, 5)
net.to(device)
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.0001)
best_acc = 0.0