[Python] When running the Tformer model training function, all label predictions evaluated are...

I tried to run the code in the Tformer paper, but when I evaluated the model in the test section of the train function, it predicted that pred_np (labels) would all be 0. My expected pred_np would be like this

[[0 12 0... 0 0 13]
[ 0 9 0 ... 14 9 0]
[14 14 14 ... 0 9 12]
[14 14 0 ... 13 0 0]
[ 0 0 10 ... 0 0 0]
[0 14 0... 12 0 9]


, I checked the differences between my train and test sections, and after deleting the same content, the remaining differences are as follows:

####################
# Train
####################
model.train()
for data, label, seg in train_loader:
sample_count += label.shape[0]
label_one_hot = label
opt.zero_grad()


loss.backward()


with torch.no_grad():
opt.step()
pred = seg_pred.max(dim=2)[1]
seg_np = seg.cpu().numpy()
pred_np = pred.detach().cpu().numpy()

####################
# Test
####################
model.eval()
with torch.no_grad():
for data, label, seg in test_loader:
sample_count += data.shape[0]
label_one_hot = label

pred = seg_pred.max(dim=2)[1]

seg_np = seg.cpu().numpy()
pred_np = pred.detach().cpu().numpy()



I suspect it's because the model hasn't updated its weights, so I tried adding code like model. load_state-dict (torch. load (savepath)) before testing, but it didn't work。Here is my complete code


def train(args, io):
save_path = args.save_path + "/exp/models/best_model.t7"

train_dataset = Teeth(num_points=args.num_points, ROOT_PATH=args.data_path, partition="train")
if len(train_dataset) < 100:
drop_last = False
else:
drop_last = True
train_loader = DataLoader(train_dataset, num_workers=8, batch_size=args.batch_size, shuffle=True,
drop_last=drop_last)
test_loader = DataLoader(Teeth(num_points=args.num_points, ROOT_PATH=args.data_path, partition="val"),
num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False)

device = torch.device("cuda:0")
print("Train dataset length", len(train_loader.dataset))
print("Val dataset length", len(test_loader.dataset))
# Try to load models
seg_num_all = 33
model = TSegFormer(args, seg_num_all).to(device)
model = nn.DataParallel(model, device_ids=[0, 1]).cuda(device)

#加载模型
if args.model_path: # 如果提供了预训练模型路径
model.load_state_dict(torch.load(args.model_path,weights_only=True))
print(f"Loaded pretrained model from {args.model_path}")

print("Let's use", torch.cuda.device_count(), "GPUs!")

if args.use_sgd:
print("Use SGD")
opt = optim.SGD(model.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=1e-4)
else:
print("Use Adam")
opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4)

if args.scheduler == 'cos':
scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=1e-3)
elif args.scheduler == 'step':
scheduler = StepLR(opt, step_size=20, gamma=0.5)

criterion1 = cal_loss
criterion2 = mask_cur_loss

best_test_iou = 0
best_test_epoch = -1
for epoch in tqdm.tqdm(range(args.epochs)):
####################
# Train
####################
train_loss = 0.0
train_loss_ging = 0.0
train_loss_seg = 0.0
train_loss_cur = 0.0
count = 0.0
model.train()
train_true_cls = []
train_pred_cls = []
train_true_seg = []
train_pred_seg = []
train_label_seg = []
sample_count = 0
for data, label, seg in train_loader:
sample_count += label.shape[0]
label_one_hot = label

AvgAngle_curvature = data[:, :, 7]
high_cur_num = int(len(AvgAngle_curvature[0]) * args.cur_threshold)
AvgAngle_idx = np.argpartition(AvgAngle_curvature, -high_cur_num, axis=1)[:, -high_cur_num:]
idx_base = np.arange(0, AvgAngle_curvature.size(0)).reshape(-1, 1) * AvgAngle_curvature.size(1)
AvgAngle_idx = AvgAngle_idx + idx_base
AvgAngle_idx = AvgAngle_idx.reshape(-1)

data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device)
AvgAngle_idx = AvgAngle_idx.to(device)
data = data.permute(0, 2, 1) # [B, 8, N]
batch_size = data.size()[0]
opt.zero_grad()

seg_pred, ging_pred = model(data, label_one_hot) # [B, 33, N], [B, 2, N]
seg_pred = seg_pred.permute(0, 2, 1).contiguous()
ging_pred = ging_pred.permute(0, 2, 1).contiguous()
ging_gold = (seg.view(-1, 1).squeeze() > 0).to(dtype=torch.int64)

loss_ging = args.w_ging * criterion1(ging_pred.view(-1, 2), ging_gold)
loss_seg = criterion1(seg_pred.view(-1, seg_num_all), seg.view(-1, 1).squeeze())
loss_cur = args.w_cur * criterion2(seg_pred.view(-1, seg_num_all), seg.view(-1, 1).squeeze(), AvgAngle_idx)
loss = loss_ging + loss_seg + loss_cur
loss.backward()
with torch.no_grad():
opt.step()
pred = seg_pred.max(dim=2)[1]
count += batch_size
train_loss += loss.item() * batch_size
train_loss_ging += loss_ging.item() * batch_size
train_loss_cur += loss_cur.item() * batch_size
train_loss_seg += loss_seg.item() * batch_size
seg_np = seg.cpu().numpy()
pred_np = pred.detach().cpu().numpy()
print(f"epoch{epoch} the train predict is {pred_np}")
train_true_cls.append(seg_np.reshape(-1))
train_pred_cls.append(pred_np.reshape(-1))
train_true_seg.append(seg_np)
train_pred_seg.append(pred_np)
train_label_seg.append(label[:, 0].type(torch.int32).reshape(-1))
if args.scheduler == 'cos':
scheduler.step()
elif args.scheduler == 'step':
if opt.param_groups[0]['lr'] > 1e-5:
scheduler.step()
if opt.param_groups[0]['lr'] < 1e-5:
for param_group in opt.param_groups:
param_group['lr'] = 1e-5
train_true_cls = np.concatenate(train_true_cls)
train_pred_cls = np.concatenate(train_pred_cls)
train_acc = metrics.accuracy_score(train_true_cls, train_pred_cls)
avg_per_class_acc = metrics.balanced_accuracy_score(train_true_cls, train_pred_cls)

train_true_seg = np.concatenate(train_true_seg, axis=0)
train_pred_seg = np.concatenate(train_pred_seg, axis=0)
train_label_seg = np.concatenate(train_label_seg)
train_ious = calculate_shape_IoU(train_pred_seg, train_true_seg, train_label_seg)
outstr = 'Train %d, loss: %.6f, loss_ging: %.6f, loss_seg: %.6f, loss_cur: %.6f, ' \
'train acc: %.6f, train avg acc: %.6f, train iou: %.6f' % \
(epoch,
train_loss * 1.0 / count,
train_loss_ging * 1.0 / count,
train_loss_seg * 1.0 / count,
train_loss_cur * 1.0 / count,
train_acc,
avg_per_class_acc,
np.mean(train_ious))
io.cprint(outstr)

####################
# Test
####################
test_loss = 0.0
test_loss_ging = 0.0
test_loss_seg = 0.0
test_loss_cur = 0.0
count = 0.0
model.eval()
test_true_cls = []
test_pred_cls = []
test_true_seg = []
test_pred_seg = []
test_label_seg = []
sample_count = 0
with torch.no_grad():
for data, label, seg in test_loader:
sample_count += data.shape[0]
label_one_hot = label

AvgAngle_curvature = data[:, :, 7]
high_cur_num = int(len(AvgAngle_curvature[0]) * args.cur_threshold)
AvgAngle_idx = np.argpartition(AvgAngle_curvature, -high_cur_num, axis=1)[:, -high_cur_num:]
idx_base = np.arange(0, AvgAngle_curvature.size(0)).reshape(-1, 1) * AvgAngle_curvature.size(1)
AvgAngle_idx = AvgAngle_idx + idx_base
AvgAngle_idx = AvgAngle_idx.reshape(-1)

data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device)
AvgAngle_idx = AvgAngle_idx.to(device)
data = data.permute(0, 2, 1) # [B, 8, N]
batch_size = data.size()[0]

seg_pred, ging_pred = model(data, label_one_hot) # [B, 33, N], [B, 2, N]
seg_pred = seg_pred.permute(0, 2, 1).contiguous()
ging_pred = ging_pred.permute(0, 2, 1).contiguous()
ging_gold = (seg.view(-1, 1).squeeze() > 0).to(dtype=torch.int64)

loss_ging = args.w_ging * criterion1(ging_pred.view(-1, 2), ging_gold)
loss_seg = criterion1(seg_pred.view(-1, seg_num_all), seg.view(-1, 1).squeeze())
loss_cur = args.w_cur * criterion2(seg_pred.view(-1, seg_num_all), seg.view(-1, 1).squeeze(),
AvgAngle_idx)
loss = loss_ging + loss_seg + loss_cur

pred = seg_pred.max(dim=2)[1]
count += batch_size
test_loss += loss.item() * batch_size
test_loss_ging += loss_ging.item() * batch_size
test_loss_cur += loss_cur.item() * batch_size
test_loss_seg += loss_seg.item() * batch_size
seg_np = seg.cpu().numpy()
pred_np = pred.detach().cpu().numpy()
print(f"epoch{epoch} the test predict is {pred_np}")
test_true_cls.append(seg_np.reshape(-1))
test_pred_cls.append(pred_np.reshape(-1))
test_true_seg.append(seg_np)
test_pred_seg.append(pred_np)
test_label_seg.append(label[:, 0].type(torch.int32).reshape(-1))
test_true_cls = np.concatenate(test_true_cls)
test_pred_cls = np.concatenate(test_pred_cls)
test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls)
avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls)
test_true_seg = np.concatenate(test_true_seg, axis=0)
test_pred_seg = np.concatenate(test_pred_seg, axis=0)
test_label_seg = np.concatenate(test_label_seg)
test_ious = calculate_shape_IoU(test_pred_seg, test_true_seg, test_label_seg)
outstr = 'Test :: loss: %.6f, test_loss_ging: %.6f, test_loss_seg: %.6f, test_loss_cur: %.6f, ' \
'test acc: %.6f, test avg acc: %.6f, test iou: %.6f, ' % \
(test_loss * 1.0 / count,
test_loss_ging * 1.0 / count,
test_loss_seg * 1.0 / count,
test_loss_cur * 1.0 / count,
test_acc,
avg_per_class_acc,
np.mean(test_ious))
io.cprint(outstr)
if np.mean(test_ious) >= best_test_iou:
best_test_iou = np.mean(test_ious)
best_test_epoch = epoch
# torch.save(model.state_dict(), os.path.join(args.save_path, '/%s/models/best_model.t7' % args.exp_name))
# save_path=os.path.join(args.save_path, '/models/best_model.t7')

torch.save(model.state_dict(), save_path)
io.cprint("best_test_iou: " + str(best_test_iou) + '; best_test_epoch: ' + str(best_test_epoch))





import os
import sys
import copy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from util import sample_and_group


class Local_op(nn.Module):
def __init__(self, in_channels, out_channels):
super(Local_op, self).__init__()
self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm1d(out_channels)
self.bn2 = nn.BatchNorm1d(out_channels)

def forward(self, x):
b, n, s, d = x.size()
x = x.permute(0, 1, 3, 2)
x = x.reshape(-1, d, s)
BtimesN, _, k = x.size()
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.adaptive_max_pool1d(x, 1).view(BtimesN,
-1)
x = x.reshape(b, n, -1).permute(0, 2, 1)
return x


class SA_Layer(nn.Module):
def __init__(self, channels):
super(SA_Layer, self).__init__()
self.q_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
self.k_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
self.q_conv.weight = self.k_conv.weight
self.q_conv.bias = self.k_conv.bias
self.v_conv = nn.Conv1d(channels, channels, 1)
self.trans_conv = nn.Conv1d(channels, channels, 1)
self.after_norm = nn.BatchNorm1d(channels)
self.act = nn.ReLU()
self.softmax = nn.Softmax(dim=-1)

def forward(self, x):
x_q = self.q_conv(x).permute(0, 2, 1)
x_k = self.k_conv(x)
x_v = self.v_conv(x)
energy = torch.bmm(x_q, x_k)
attention = F.softmax(energy / np.sqrt(x_q.size(-1)), dim=-2)
x_r = torch.bmm(x_v, attention)
x_r = self.act(self.after_norm(self.trans_conv(x_r)))
x = x + x_r

return x


class Point_Transformer_Last(nn.Module):
def __init__(self, args, channels=128):
super(Point_Transformer_Last, self).__init__()
self.args = args
self.sa1 = SA_Layer(channels)
self.sa2 = SA_Layer(channels)
self.sa3 = SA_Layer(channels)
self.sa4 = SA_Layer(channels)

def forward(self, x):
batch_size, _, N = x.size()
x1 = self.sa1(x)
x2 = self.sa2(x1)
x3 = self.sa3(x2)
x4 = self.sa4(x3)
x = torch.cat((x1, x2, x3, x4), dim=1)

return x


class TSegFormer(nn.Module):
def __init__(self, args, part_num=33):
super(TSegFormer, self).__init__()
self.part_num = part_num
self.args = args

self.conv1 = nn.Conv1d(8, 64, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(64, 64, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(64)
self.gather_local_0 = Local_op(in_channels=128, out_channels=64)
self.gather_local_1 = Local_op(in_channels=128, out_channels=128)
self.pt_last = Point_Transformer_Last(args)
self.conv_fuse = nn.Sequential(nn.Conv1d(512, 1024, kernel_size=1, bias=False),
nn.BatchNorm1d(1024),
nn.LeakyReLU(negative_slope=0.2))
self.label_conv = nn.Sequential(nn.Conv1d(2, 64, kernel_size=1, bias=False),
nn.BatchNorm1d(64),
nn.LeakyReLU(negative_slope=0.2))

self.convs1 = nn.Conv1d(1024 * 3 + 64, 512, 1)
self.dp1 = nn.Dropout(0.5)
self.convs2 = nn.Conv1d(512, 256, 1)
self.convs3 = nn.Conv1d(256, self.part_num, 1)
self.bns1 = nn.BatchNorm1d(512)
self.bns2 = nn.BatchNorm1d(256)

self.convs4 = nn.Conv1d(1024 * 3 + 64, 512, 1)
self.dp4 = nn.Dropout(0.5)
self.convs5 = nn.Conv1d(512, 256, 1)
self.convs6 = nn.Conv1d(256, 2, 1)
self.bns4 = nn.BatchNorm1d(512)
self.bns5 = nn.BatchNorm1d(256)

def forward(self, x, cls_label):
xyz = x.permute(0, 2, 1)[:, :, 0:3]
batch_size, _, N = x.size()

x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = x.permute(0, 2, 1)
new_xyz, new_feature = sample_and_group(npoint=N, radius=0.15, nsample=32, xyz=xyz,
points=x)
feature_0 = self.gather_local_0(new_feature)
feature = feature_0.permute(0, 2, 1)
new_xyz, new_feature = sample_and_group(npoint=N, radius=0.2, nsample=32, xyz=new_xyz,
points=feature)
feature_1 = self.gather_local_1(new_feature)

x = self.pt_last(feature_1)
x = self.conv_fuse(x)

x_max_feature = F.adaptive_max_pool1d(x, 1).repeat(1, 1, N)
x_avg_feature = F.adaptive_avg_pool1d(x, 1).repeat(1, 1, N)

cls_label_one_hot = cls_label.view(batch_size, 2, 1)
# cls_label=cls_label[:,0:1,:]
cls_label_one_hot = cls_label.view(batch_size, 2, 1).float()
cls_label_feature = self.label_conv(cls_label_one_hot).repeat(1, 1, N)
x_global_feature = torch.cat([x_max_feature, x_avg_feature, cls_label_feature], 1)
all_fea = torch.cat((x, x_global_feature), 1)

x = F.relu(self.bns1(self.convs1(all_fea)))
x = self.dp1(x)
x = F.relu(self.bns2(self.convs2(x)))
x = self.convs3(x)

y = F.relu(self.bns4(self.convs4(all_fea)))
y = self.dp4(y)
y = F.relu(self.bns5(self.convs5(y)))
y = self.convs6(y)

return x, y


In summary, the previous section of the code can correctly predict the seg_pred

####################
# Train
####################
train_loss = 0.0
train_loss_ging = 0.0
train_loss_seg = 0.0
train_loss_cur = 0.0
count = 0.0
model.train()
train_true_cls = []
train_pred_cls = []
train_true_seg = []
train_pred_seg = []
train_label_seg = []
sample_count = 0
for data, label, seg in train_loader:
sample_count += label.shape[0]
label_one_hot = label

AvgAngle_curvature = data[:, :, 7]
high_cur_num = int(len(AvgAngle_curvature[0]) * args.cur_threshold)
AvgAngle_idx = np.argpartition(AvgAngle_curvature, -high_cur_num, axis=1)[:, -high_cur_num:]
idx_base = np.arange(0, AvgAngle_curvature.size(0)).reshape(-1, 1) * AvgAngle_curvature.size(1)
AvgAngle_idx = AvgAngle_idx + idx_base
AvgAngle_idx = AvgAngle_idx.reshape(-1)

data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device)
AvgAngle_idx = AvgAngle_idx.to(device)
data = data.permute(0, 2, 1) # [B, 8, N]
batch_size = data.size()[0]
opt.zero_grad()

seg_pred, ging_pred = model(data, label_one_hot) # [B, 33, N], [B, 2, N]
seg_pred = seg_pred.permute(0, 2, 1).contiguous()
ging_pred = ging_pred.permute(0, 2, 1).contiguous()
ging_gold = (seg.view(-1, 1).squeeze() > 0).to(dtype=torch.int64)

loss_ging = args.w_ging * criterion1(ging_pred.view(-1, 2), ging_gold)
loss_seg = criterion1(seg_pred.view(-1, seg_num_all), seg.view(-1, 1).squeeze())
loss_cur = args.w_cur * criterion2(seg_pred.view(-1, seg_num_all), seg.view(-1, 1).squeeze(), AvgAngle_idx)
loss = loss_ging + loss_seg + loss_cur
loss.backward()
with torch.no_grad():
opt.step()
pred = seg_pred.max(dim=2)[1]
count += batch_size
train_loss += loss.item() * batch_size
train_loss_ging += loss_ging.item() * batch_size
train_loss_cur += loss_cur.item() * batch_size
train_loss_seg += loss_seg.item() * batch_size
seg_np = seg.cpu().numpy()
pred_np = pred.detach().cpu().numpy()
print(f"epoch{epoch} the train predict is {pred_np}")
train_true_cls.append(seg_np.reshape(-1))
train_pred_cls.append(pred_np.reshape(-1))
train_true_seg.append(seg_np)
train_pred_seg.append(pred_np)
train_label_seg.append(label[:, 0].type(torch.int32).reshape(-1))
if args.scheduler == 'cos':
scheduler.step()
elif args.scheduler == 'step':
if opt.param_groups[0]['lr'] > 1e-5:
scheduler.step()
if opt.param_groups[0]['lr'] < 1e-5:
for param_group in opt.param_groups:
param_group['lr'] = 1e-5
train_true_cls = np.concatenate(train_true_cls)
train_pred_cls = np.concatenate(train_pred_cls)
train_acc = metrics.accuracy_score(train_true_cls, train_pred_cls)
avg_per_class_acc = metrics.balanced_accuracy_score(train_true_cls, train_pred_cls)

train_true_seg = np.concatenate(train_true_seg, axis=0)
train_pred_seg = np.concatenate(train_pred_seg, axis=0)
train_label_seg = np.concatenate(train_label_seg)
train_ious = calculate_shape_IoU(train_pred_seg, train_true_seg, train_label_seg)
outstr = 'Train %d, loss: %.6f, loss_ging: %.6f, loss_seg: %.6f, loss_cur: %.6f, ' \
'train acc: %.6f, train avg acc: %.6f, train iou: %.6f' % \
(epoch,
train_loss * 1.0 / count,
train_loss_ging * 1.0 / count,
train_loss_seg * 1.0 / count,
train_loss_cur * 1.0 / count,
train_acc,
avg_per_class_acc,
np.mean(train_ious))
io.cprint(outstr)


, while the following section is not

####################
# Test
####################
test_loss = 0.0
test_loss_ging = 0.0
test_loss_seg = 0.0
test_loss_cur = 0.0
count = 0.0
model.eval()
test_true_cls = []
test_pred_cls = []
test_true_seg = []
test_pred_seg = []
test_label_seg = []
sample_count = 0
with torch.no_grad():
for data, label, seg in test_loader:
sample_count += data.shape[0]
label_one_hot = label

AvgAngle_curvature = data[:, :, 7]
high_cur_num = int(len(AvgAngle_curvature[0]) * args.cur_threshold)
AvgAngle_idx = np.argpartition(AvgAngle_curvature, -high_cur_num, axis=1)[:, -high_cur_num:]
idx_base = np.arange(0, AvgAngle_curvature.size(0)).reshape(-1, 1) * AvgAngle_curvature.size(1)
AvgAngle_idx = AvgAngle_idx + idx_base
AvgAngle_idx = AvgAngle_idx.reshape(-1)

data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device)
AvgAngle_idx = AvgAngle_idx.to(device)
data = data.permute(0, 2, 1) # [B, 8, N]
batch_size = data.size()[0]

seg_pred, ging_pred = model(data, label_one_hot) # [B, 33, N], [B, 2, N]
seg_pred = seg_pred.permute(0, 2, 1).contiguous()
ging_pred = ging_pred.permute(0, 2, 1).contiguous()
ging_gold = (seg.view(-1, 1).squeeze() > 0).to(dtype=torch.int64)

loss_ging = args.w_ging * criterion1(ging_pred.view(-1, 2), ging_gold)
loss_seg = criterion1(seg_pred.view(-1, seg_num_all), seg.view(-1, 1).squeeze())
loss_cur = args.w_cur * criterion2(seg_pred.view(-1, seg_num_all), seg.view(-1, 1).squeeze(),
AvgAngle_idx)
loss = loss_ging + loss_seg + loss_cur

pred = seg_pred.max(dim=2)[1]
count += batch_size
test_loss += loss.item() * batch_size
test_loss_ging += loss_ging.item() * batch_size
test_loss_cur += loss_cur.item() * batch_size
test_loss_seg += loss_seg.item() * batch_size
seg_np = seg.cpu().numpy()
pred_np = pred.detach().cpu().numpy()
print(f"epoch{epoch} the test predict is {pred_np}")
test_true_cls.append(seg_np.reshape(-1))
test_pred_cls.append(pred_np.reshape(-1))
test_true_seg.append(seg_np)
test_pred_seg.append(pred_np)
test_label_seg.append(label[:, 0].type(torch.int32).reshape(-1))
test_true_cls = np.concatenate(test_true_cls)
test_pred_cls = np.concatenate(test_pred_cls)
test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls)
avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls)
test_true_seg = np.concatenate(test_true_seg, axis=0)
test_pred_seg = np.concatenate(test_pred_seg, axis=0)
test_label_seg = np.concatenate(test_label_seg)
test_ious = calculate_shape_IoU(test_pred_seg, test_true_seg, test_label_seg)
outstr = 'Test :: loss: %.6f, test_loss_ging: %.6f, test_loss_seg: %.6f, test_loss_cur: %.6f, ' \
'test acc: %.6f, test avg acc: %.6f, test iou: %.6f, ' % \
(test_loss * 1.0 / count,
test_loss_ging * 1.0 / count,
test_loss_seg * 1.0 / count,
test_loss_cur * 1.0 / count,
test_acc,
avg_per_class_acc,
np.mean(test_ious))
io.cprint(outstr)
if np.mean(test_ious) >= best_test_iou:
best_test_iou = np.mean(test_ious)
best_test_epoch = epoch
# torch.save(model.state_dict(), os.path.join(args.save_path, '/%s/models/best_model.t7' % args.exp_name))
# save_path=os.path.join(args.save_path, '/models/best_model.t7')

torch.save(model.state_dict(), save_path)
io.cprint("best_test_iou: " + str(best_test_iou) + '; best_test_epoch: ' + str(best_test_epoch))

Continue reading...