def CreatNet(opt):
name = opt.model_name
label_num = opt.label
if name == 'lstm':
net = lstm.lstm(opt.input_size,
opt.time_step,
input_nc=opt.input_nc,
num_classes=label_num)
elif name == 'cnn_1d':
net = cnn_1d.cnn(opt.input_nc, num_classes=label_num)
elif name == 'resnet18_1d':
net = resnet_1d.resnet18()
net.conv1 = nn.Conv1d(opt.input_nc, 64, 7, 2, 3, bias=False)
net.fc = nn.Linear(512, label_num)
elif name == 'resnet34_1d':
net = resnet_1d.resnet34()
net.conv1 = nn.Conv1d(opt.input_nc, 64, 7, 2, 3, bias=False)
net.fc = nn.Linear(512, label_num)
elif name == 'multi_scale_resnet_1d':
net = multi_scale_resnet_1d.Multi_Scale_ResNet(inchannel=opt.input_nc,
num_classes=label_num)
elif name == 'micro_multi_scale_resnet_1d':
net = micro_multi_scale_resnet_1d.Multi_Scale_ResNet(
inchannel=opt.input_nc, num_classes=label_num)
elif name == 'multi_scale_resnet':
net = multi_scale_resnet.Multi_Scale_ResNet(inchannel=opt.input_nc,
num_classes=label_num)
elif name == 'dfcnn':
net = dfcnn.dfcnn(num_classes=label_num)
elif name in ['resnet101', 'resnet50', 'resnet18']:
if name == 'resnet101':
net = resnet.resnet101(pretrained=False)
net.fc = nn.Linear(2048, label_num)
elif name == 'resnet50':
net = resnet.resnet50(pretrained=False)
net.fc = nn.Linear(2048, label_num)
elif name == 'resnet18':
net = resnet.resnet18(pretrained=False)
net.fc = nn.Linear(512, label_num)
net.conv1 = nn.Conv2d(opt.input_nc, 64, 7, 2, 3, bias=False)
elif 'densenet' in name:
if name == 'densenet121':
net = densenet.densenet121(pretrained=False, num_classes=label_num)
elif name == 'densenet201':
net = densenet.densenet201(pretrained=False, num_classes=label_num)
elif name == 'squeezenet':
net = squeezenet.squeezenet1_1(pretrained=False,
num_classes=label_num,
inchannel=1)
return net
def CreatNet(name):
if name == 'lstm':
net = lstm.lstm(100, 27, num_classes=5)
elif name == 'cnn_1d':
net = cnn_1d.cnn(1, num_classes=5)
elif name == 'resnet18_1d':
net = resnet_1d.resnet18()
net.conv1 = nn.Conv1d(1, 64, 7, 2, 3, bias=False)
net.fc = nn.Linear(512, 5)
elif name == 'multi_scale_resnet_1d':
net = multi_scale_resnet_1d.Multi_Scale_ResNet(inchannel=1,
num_classes=5)
elif name == 'multi_scale_resnet':
net = multi_scale_resnet.Multi_Scale_ResNet(inchannel=1, num_classes=5)
elif name == 'dfcnn':
net = dfcnn.dfcnn(num_classes=5)
elif name in ['resnet101', 'resnet50', 'resnet18']:
if name == 'resnet101':
net = resnet.resnet101(pretrained=False)
net.fc = nn.Linear(2048, 5)
elif name == 'resnet50':
net = resnet.resnet50(pretrained=False)
net.fc = nn.Linear(2048, 5)
elif name == 'resnet18':
net = resnet.resnet18(pretrained=False)
net.fc = nn.Linear(512, 5)
net.conv1 = nn.Conv2d(1, 64, 7, 2, 3, bias=False)
elif 'densenet' in name:
if name == 'densenet121':
net = densenet.densenet121(pretrained=False, num_classes=5)
elif name == 'densenet201':
net = densenet.densenet201(pretrained=False, num_classes=5)
elif name == 'squeezenet':
net = squeezenet.squeezenet1_1(pretrained=False,
num_classes=5,
inchannel=1)
return net
def train_model(modname='alexnet', pm_ch='both', bs=16):
"""
Args:
modname (string): Name of the model. Has to be one of the values:
'alexnet', batch 64
'densenet'
'inception'
'resnet', batch 16
'squeezenet', batch 16
'vgg'
pm_ch (string): pixelmap channel -- 'time', 'charge', 'both', default to both
"""
# device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# hyper parameters
max_epochs = 10
learning_rate = 0.001
# determine number of input channels
nch = 2
if pm_ch != 'both':
nch = 1
ds = PixelMapDataset('training_file_list.txt', pm_ch)
# try out the data loader utility
dl = torch.utils.data.DataLoader(dataset=ds, batch_size=bs, shuffle=True)
# define model
model = None
if modname == 'alexnet':
model = alexnet(num_classes=3, in_ch=nch).to(device)
elif modname == 'densenet':
model = DenseNet(num_classes=3, in_ch=nch).to(device)
elif modname == 'inception':
model = inception_v3(num_classes=3, in_ch=nch).to(device)
elif modname == 'resnet':
model = resnet18(num_classes=3, in_ch=nch).to(device)
elif modname == 'squeezenet':
model = squeezenet1_1(num_classes=3, in_ch=nch).to(device)
elif modname == 'vgg':
model = vgg19_bn(in_ch=nch, num_classes=3).to(device)
else:
print('Model {} not defined.'.format(modname))
return
# loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# training process
total_step = len(dl)
for epoch in range(max_epochs):
for i, (view1, view2, local_labels) in enumerate(dl):
view1 = view1.float().to(device)
if modname == 'inception':
view1 = nn.ZeroPad2d((0, 192, 102, 101))(view1)
else:
view1 = nn.ZeroPad2d((0, 117, 64, 64))(view1)
local_labels = local_labels.to(device)
# forward pass
outputs = model(view1)
loss = criterion(outputs, local_labels)
# backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % bs == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
epoch + 1, max_epochs, i + 1, total_step, loss.item()))
# save the model checkpoint
save_path = '../../../data/two_views/saved_models/{}/{}'.format(
modname, pm_ch)
os.makedirs(save_path, exist_ok=True)
torch.save(model.state_dict(), os.path.join(save_path, 'model.ckpt'))
def get_model(class_num):
if (MODEL_TYPE == 'alexnet'):
model = alexnet.alexnet(pretrained=FINETUNE)
elif (MODEL_TYPE == 'vgg'):
if (MODEL_DEPTH_OR_VERSION == 11):
model = vgg.vgg11(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 13):
model = vgg.vgg13(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 16):
model = vgg.vgg16(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 19):
model = vgg.vgg19(pretrained=FINETUNE)
else:
print('Error : VGG should have depth of either [11, 13, 16, 19]')
sys.exit(1)
elif (MODEL_TYPE == 'squeezenet'):
if (MODEL_DEPTH_OR_VERSION == 0 or MODEL_DEPTH_OR_VERSION == 'v0'):
model = squeezenet.squeezenet1_0(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 1 or MODEL_DEPTH_OR_VERSION == 'v1'):
model = squeezenet.squeezenet1_1(pretrained=FINETUNE)
else:
print('Error : Squeezenet should have version of either [0, 1]')
sys.exit(1)
elif (MODEL_TYPE == 'resnet'):
if (MODEL_DEPTH_OR_VERSION == 18):
model = resnet.resnet18(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 34):
model = resnet.resnet34(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 50):
model = resnet.resnet50(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 101):
model = resnet.resnet101(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 152):
model = resnet.resnet152(pretrained=FINETUNE)
else:
print(
'Error : Resnet should have depth of either [18, 34, 50, 101, 152]'
)
sys.exit(1)
elif (MODEL_TYPE == 'densenet'):
if (MODEL_DEPTH_OR_VERSION == 121):
model = densenet.densenet121(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 169):
model = densenet.densenet169(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 161):
model = densenet.densenet161(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 201):
model = densenet.densenet201(pretrained=FINETUNE)
else:
print(
'Error : Densenet should have depth of either [121, 169, 161, 201]'
)
sys.exit(1)
elif (MODEL_TYPE == 'inception'):
if (MODEL_DEPTH_OR_VERSION == 3 or MODEL_DEPTH_OR_VERSION == 'v3'):
model = inception.inception_v3(pretrained=FINETUNE)
else:
print('Error : Inception should have version of either [3, ]')
sys.exit(1)
else:
print(
'Error : Network should be either [alexnet / squeezenet / vgg / resnet / densenet / inception]'
)
sys.exit(1)
if (MODEL_TYPE == 'alexnet' or MODEL_TYPE == 'vgg'):
num_ftrs = model.classifier[6].in_features
feature_model = list(model.classifier.children())
feature_model.pop()
feature_model.append(nn.Linear(num_ftrs, class_num))
model.classifier = nn.Sequential(*feature_model)
elif (MODEL_TYPE == 'resnet' or MODEL_TYPE == 'inception'):
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, class_num)
elif (MODEL_TYPE == 'densenet'):
num_ftrs = model.classifier.in_features
model.classifier = nn.Linear(num_ftrs, class_num)
return model
batch_size=batch_size,
shuffle=True,
num_workers=4)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_name = args.backbone
if model_name == "vgg16":
model = vgg16(architecture_type="cam",
pretrained=True,
num_classes=n_classes,
large_feature_map=False)
elif model_name == "resnet18":
model = resnet18('cam', num_classes=n_classes)
elif model_name == "squeezenet1_1":
model = squeezenet1_1(num_classes=n_classes)
print(model)
checkpoint = torch.load(args.checkpoint_path, map_location=torch.device('cpu'))
model.load_state_dict(checkpoint)
root = args.dataset_path + 'val/'
masks_root = args.masks_path
import glob
filename = [
glob.glob(root + i + "/*.jpg") for i in object_categories
if i not in ["046_plastic_bolt"]
]
filename_list = [item for sublist in filename for item in sublist]
# print(filename_list)
count = 0
if ops.model == 'resnet_50':
model_ = resnet50(num_classes=ops.num_classes,
img_size=ops.img_size[0])
elif ops.model == 'resnet_18':
model_ = resnet18(num_classes=ops.num_classes,
img_size=ops.img_size[0])
elif ops.model == 'resnet_34':
model_ = resnet34(num_classes=ops.num_classes,
img_size=ops.img_size[0])
elif ops.model == 'resnet_101':
model_ = resnet101(num_classes=ops.num_classes,
img_size=ops.img_size[0])
elif ops.model == "squeezenet1_0":
model_ = squeezenet1_0(num_classes=ops.num_classes)
elif ops.model == "squeezenet1_1":
model_ = squeezenet1_1(num_classes=ops.num_classes)
elif ops.model == "shufflenetv2":
model_ = ShuffleNetV2(ratio=1., num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x1_5":
model_ = shufflenet_v2_x1_5(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x1_0":
model_ = shufflenet_v2_x1_0(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x2_0":
model_ = shufflenet_v2_x2_0(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet":
model_ = ShuffleNet(num_blocks=[2, 4, 2],
num_classes=ops.num_classes,
groups=3)
def test_model(modname='alexnet', pm_ch='both', bs=16):
# hyperparameters
batch_size = bs
# device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# determine number of input channels
nch = 2
if pm_ch != 'both':
nch = 1
# restore model
model = None
if modname == 'alexnet':
model = alexnet(num_classes=3, in_ch=nch).to(device)
elif modname == 'densenet':
model = DenseNet(num_classes=3, in_ch=nch).to(device)
elif modname == 'inception':
model = inception_v3(num_classes=3, in_ch=nch).to(device)
elif modname == 'resnet':
model = resnet18(num_classes=3, in_ch=nch).to(device)
elif modname == 'squeezenet':
model = squeezenet1_1(num_classes=3, in_ch=nch).to(device)
elif modname == 'vgg':
model = vgg19_bn(in_ch=nch, num_classes=3).to(device)
else:
print('Model {} not defined.'.format(modname))
return
# retrieve trained model
# load path
load_path = '../../../data/two_views/saved_models/{}/{}'.format(
modname, pm_ch)
model_pathname = os.path.join(load_path, 'model.ckpt')
if not os.path.exists(model_pathname):
print('Trained model file {} does not exist. Abort.'.format(
model_pathname))
return
model.load_state_dict(torch.load(model_pathname))
# load test dataset
test_dataset = PixelMapDataset('test_file_list.txt', pm_ch)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
# test the model
model.eval(
) # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
with torch.no_grad():
correct = 0
total = 0
correct_cc_or_bkg = 0
ws_total = 0
ws_correct = 0
for view1, view2, labels in test_loader:
view1 = view1.float().to(device)
if modname == 'inception':
view1 = nn.ZeroPad2d((0, 192, 102, 101))(view1)
else:
view1 = nn.ZeroPad2d((0, 117, 64, 64))(view1)
labels = labels.to(device)
outputs = model(view1)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
for i in range(len(predicted)):
if (predicted[i] < 2
and labels[i] < 2) or (predicted[i] == 2
and labels[i] == 2):
correct_cc_or_bkg += 1
if labels[i] < 2:
ws_total += 1
if (predicted[i] == labels[i]):
ws_correct += 1
print('Model Performance:')
print('Model:', modname)
print('Channel:', pm_ch)
print(
'3-class Test Accuracy of the model on the test images: {}/{}, {:.2f} %'
.format(correct, total, 100 * correct / total))
print(
'2-class Test Accuracy of the model on the test images: {}/{}, {:.2f} %'
.format(correct_cc_or_bkg, total, 100 * correct_cc_or_bkg / total))
print(
'Wrong-sign Test Accuracy of the model on the test images: {}/{}, {:.2f} %'
.format(ws_correct, ws_total, 100 * ws_correct / ws_total))
def trainer(ops, f_log):
try:
os.environ['CUDA_VISIBLE_DEVICES'] = ops.GPUS
if ops.log_flag:
sys.stdout = f_log
set_seed(ops.seed)
#---------------------------------------------------------------- 构建模型
if ops.model == 'resnet_50':
model_ = resnet50(pretrained=True,
num_classes=ops.num_classes,
img_size=ops.img_size[0],
dropout_factor=ops.dropout)
elif ops.model == 'resnet_18':
model_ = resnet18(pretrained=True,
num_classes=ops.num_classes,
img_size=ops.img_size[0],
dropout_factor=ops.dropout)
elif ops.model == 'resnet_34':
model_ = resnet34(pretrained=True,
num_classes=ops.num_classes,
img_size=ops.img_size[0],
dropout_factor=ops.dropout)
elif ops.model == 'resnet_101':
model_ = resnet101(pretrained=True,
num_classes=ops.num_classes,
img_size=ops.img_size[0],
dropout_factor=ops.dropout)
elif ops.model == "squeezenet1_0":
model_ = squeezenet1_0(pretrained=True,
num_classes=ops.num_classes,
dropout_factor=ops.dropout)
elif ops.model == "squeezenet1_1":
model_ = squeezenet1_1(pretrained=True,
num_classes=ops.num_classes,
dropout_factor=ops.dropout)
elif ops.model == "shufflenetv2":
model_ = ShuffleNetV2(ratio=1.,
num_classes=ops.num_classes,
dropout_factor=ops.dropout)
elif ops.model == "shufflenet_v2_x1_5":
model_ = shufflenet_v2_x1_5(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x1_0":
model_ = shufflenet_v2_x1_0(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x2_0":
model_ = shufflenet_v2_x2_0(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet":
model_ = ShuffleNet(num_blocks=[2, 4, 2],
num_classes=ops.num_classes,
groups=3,
dropout_factor=ops.dropout)
elif ops.model == "mobilenetv2":
model_ = MobileNetV2(num_classes=ops.num_classes,
dropout_factor=ops.dropout)
elif ops.model == "ReXNetV1":
model_ = ReXNetV1(num_classes=ops.num_classes,
dropout_factor=ops.dropout)
else:
print(" no support the model")
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
model_ = model_.to(device)
# print(model_)# 打印模型结构
# Dataset
dataset = LoadImagesAndLabels(ops=ops,
img_size=ops.img_size,
flag_agu=ops.flag_agu,
fix_res=ops.fix_res,
vis=False)
print("handpose done")
print('len train datasets : %s' % (dataset.__len__()))
# Dataloader
dataloader = DataLoader(dataset,
batch_size=ops.batch_size,
num_workers=ops.num_workers,
shuffle=True,
pin_memory=False,
drop_last=True)
# 优化器设计
optimizer_Adam = torch.optim.Adam(model_.parameters(),
lr=ops.init_lr,
betas=(0.9, 0.99),
weight_decay=1e-6)
# optimizer_SGD = optim.SGD(model_.parameters(), lr=ops.init_lr, momentum=ops.momentum, weight_decay=ops.weight_decay)# 优化器初始化
optimizer = optimizer_Adam
# 加载 finetune 模型
if os.access(ops.fintune_model, os.F_OK): # checkpoint
chkpt = torch.load(ops.fintune_model, map_location=device)
model_.load_state_dict(chkpt)
print('load fintune model : {}'.format(ops.fintune_model))
print('/**********************************************/')
# 损失函数
if ops.loss_define != 'wing_loss':
criterion = nn.MSELoss(reduce=True, reduction='mean')
step = 0
idx = 0
# 变量初始化
best_loss = np.inf
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
flag_change_lr_cnt = 0 # 学习率更新计数器
init_lr = ops.init_lr # 学习率
epochs_loss_dict = {}
for epoch in range(0, ops.epochs):
if ops.log_flag:
sys.stdout = f_log
print('\nepoch %d ------>>>' % epoch)
model_.train()
# 学习率更新策略
if loss_mean != 0.:
if best_loss > (loss_mean / loss_idx):
flag_change_lr_cnt = 0
best_loss = (loss_mean / loss_idx)
else:
flag_change_lr_cnt += 1
if flag_change_lr_cnt > 50:
init_lr = init_lr * ops.lr_decay
set_learning_rate(optimizer, init_lr)
flag_change_lr_cnt = 0
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
for i, (imgs_, pts_) in enumerate(dataloader):
# print('imgs_, pts_',imgs_.size(), pts_.size())
if use_cuda:
imgs_ = imgs_.cuda(
) # pytorch 的 数据输入格式 ： (batch, channel, height, width)
pts_ = pts_.cuda()
output = model_(imgs_.float())
if ops.loss_define == 'wing_loss':
loss = got_total_wing_loss(output, pts_.float())
else:
loss = criterion(output, pts_.float())
loss_mean += loss.item()
loss_idx += 1.
if i % 10 == 0:
loc_time = time.strftime("%Y-%m-%d %H:%M:%S",
time.localtime())
print(' %s - %s - epoch [%s/%s] (%s/%s):'%(loc_time,ops.model,epoch,ops.epochs,i,int(dataset.__len__()/ops.batch_size)),\
'Mean Loss : %.6f - Loss: %.6f'%(loss_mean/loss_idx,loss.item()),\
' lr : %.8f'%init_lr,' bs :',ops.batch_size,\
' img_size: %s x %s'%(ops.img_size[0],ops.img_size[1]),' best_loss: %.6f'%best_loss)
# 计算梯度
loss.backward()
# 优化器对模型参数更新
optimizer.step()
# 优化器梯度清零
optimizer.zero_grad()
step += 1
torch.save(
model_.state_dict(),
ops.model_exp + '{}-size-{}-model_epoch-{}.pth'.format(
ops.model, ops.img_size[0], epoch))
except Exception as e:
print('Exception : ', e) # 打印异常
print('Exception file : ',
e.__traceback__.tb_frame.f_globals['__file__']) # 发生异常所在的文件
print('Exception line : ', e.__traceback__.tb_lineno) # 发生异常所在的行数
