def save_feature_batch(model_name, ckpt_path, feature_path, label_path, data_type="valid"):
model = create_model(model_name.split('-')[-2], num_classes=3, checkpoint_path=ckpt_path)
model.cuda().eval()
print('..... Finished loading model! ......')
img_size = int(model_name.split('-')[-1])
interpolation = "bicubic"
batch_size = 128
dataset = Dataset(os.path.join(BASE, data_type))
loader = create_loader(
dataset,
input_size=img_size,
batch_size=batch_size,
use_prefetcher=False,
interpolation=interpolation,
num_workers=8)
features = []
labels = []
with torch.no_grad():
for batch_idx, (input, target) in enumerate(tqdm(loader)):
out = model.forward_features(input.cuda())
out2 = nn.AdaptiveAvgPool2d(1)(out)
feature = out2.view(out.size(0), -1)
features.append(feature.cpu().numpy())
labels.extend(target.cpu().numpy())
features = np.array(np.vstack(features))
pickle.dump(features, open(feature_path, 'wb'))
pickle.dump(labels, open(label_path, 'wb'))
print('CNN features obtained and saved.')
def save_feature_test(model_name, ckpt_path, feature_path, imgs):
'''
提取特征，保存为pkl文件
'''
model = create_model(model_name.split('-')[-2], num_classes=3, checkpoint_path=ckpt_path)
model.eval().cuda()
print('..... Finished loading model! ......')
## 特征的维度需要自己根据特定的模型调整，我这里采用的是哪一个我也忘了
nb_features = 2048
features = np.empty((len(imgs), nb_features))
for i in tqdm(range(len(imgs))):
img_path = imgs[i].strip().split(' ')[0]
img = Image.open(img_path).convert('RGB')
# bilinear bicubic
img = get_test_transform(img_size=int(model_name.split('-')[-1]), interpolation="bicubic")(img).unsqueeze(
0).cuda()
with torch.no_grad():
out = model.forward_features(img)
out2 = nn.AdaptiveAvgPool2d(1)(out)
feature = out2.view(out.size(1), -1).squeeze(1)
features[i, :] = feature.cpu().numpy()
pickle.dump(features, open(feature_path, 'wb'))
print('CNN features obtained and saved.')
def main():
args = parser.parse_args()
args.batch_size = args.batch_size
# setup model
print('creating and loading the model...')
state = torch.load(args.model_path, map_location='cpu')
args.num_classes = state['num_classes']
args.do_bottleneck_head = False
model = create_model(args).cuda()
model.load_state_dict(state['model'], strict=True)
model.eval()
classes_list = np.array(list(state['idx_to_class'].values()))
print('done\n')
# Data loading code
normalize = transforms.Normalize(mean=[0, 0, 0],
std=[1, 1, 1])
instances_path = os.path.join(args.data, 'annotations/instances_val2017.json')
data_path = os.path.join(args.data, 'val2017')
val_dataset = CocoDetection(data_path,
instances_path,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
normalize,
]))
print("len(val_dataset)): ", len(val_dataset))
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
validate_multi(val_loader, model, args)
def main():
print('ASL Example Inference code on a single image')
# parsing args
args = parse_args(parser)
# setup model
print('creating and loading the model...')
state = torch.load(args.model_path, map_location='cpu')
args.num_classes = state['num_classes']
print('num_classes: ', args.num_classes)
classes_list = np.array(list(state['idx_to_class'].values()))
print('classes_list: ', classes_list)
model = create_model(args).cuda()
model.load_state_dict(state['model'], strict=True)
model.eval()
print('sigmoid threshold: ', args.th)
# doing inference
print('loading image and doing inference...')
im = Image.open(args.pic_path)
im_resize = im.resize((args.input_size, args.input_size))
np_img = np.array(im_resize, dtype=np.uint8)
tensor_img = torch.from_numpy(np_img).permute(
2, 0, 1).float() / 255.0 # HWC to CHW
tensor_batch = torch.unsqueeze(tensor_img, 0).cuda()
output = torch.squeeze(torch.sigmoid(model(tensor_batch)))
np_output = output.cpu().detach().numpy()
detected_classes = classes_list[np_output > args.th]
print('done\n')
# example loss calculation
output = model(tensor_batch)
loss_func1 = AsymmetricLoss()
loss_func2 = AsymmetricLossOptimized()
target = output.clone()
target[output < 0] = 0 # mockup target
target[output >= 0] = 1
loss1 = loss_func1(output, target)
loss2 = loss_func2(output, target)
assert abs((loss1.item() - loss2.item())) < 1e-6
# displaying image
print('showing image on screen...')
fig = plt.figure()
plt.imshow(im)
plt.axis('off')
plt.axis('tight')
# plt.rcParams["axes.titlesize"] = 10
plt.title("detected classes: {}".format(detected_classes))
plt.show()
print('done\n')
def main():
args = parser.parse_args()
args.do_bottleneck_head = False
# Setup model
print('creating model...')
model = create_model(args).cuda()
if args.model_path: # make sure to load pretrained ImageNet model
state = torch.load(args.model_path, map_location='cpu')
filtered_dict = {k: v for k, v in state['model'].items() if
(k in model.state_dict() and 'head.fc' not in k)}
model.load_state_dict(filtered_dict, strict=False)
print('done\n')
# COCO Data loading
instances_path_val = os.path.join(args.data, 'annotations/instances_val2014.json')
instances_path_train = os.path.join(args.data, 'annotations/instances_train2014.json')
# data_path_val = args.data
# data_path_train = args.data
data_path_val = f'{args.data}/val2014' # args.data
data_path_train = f'{args.data}/train2014' # args.data
val_dataset = CocoDetection(data_path_val,
instances_path_val,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
# normalize, # no need, toTensor does normalization
]))
train_dataset = CocoDetection(data_path_train,
instances_path_train,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
CutoutPIL(cutout_factor=0.5),
RandAugment(),
transforms.ToTensor(),
# normalize,
]))
print("len(val_dataset)): ", len(val_dataset))
print("len(train_dataset)): ", len(train_dataset))
# Pytorch Data loader
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=False)
# Actuall Training
train_multi_label_coco(model, train_loader, val_loader, args.lr)
def main():
print('ASL Example Inference code on a single image')
# parsing args
args = parse_args(parser)
# setup model
print('creating and loading the model...')
state = torch.load(args.model_path, map_location='cpu')
args.num_classes = state['num_classes']
model = create_model(args).cuda()
model.load_state_dict(state['model'], strict=True)
model.eval()
classes_list = np.array(list(state['idx_to_class'].values()))
print('done\n')
print('loading image and doing inference...')
im = Image.open(args.pic_path)
im_resize = im.resize((args.input_size, args.input_size))
np_img = np.array(im_resize, dtype=np.uint8)
tensor_img = torch.from_numpy(np_img).permute(
2, 0, 1).float() / 255.0 # HWC to CHW
print(tensor_img.shape)
tensor_batch = torch.unsqueeze(tensor_img, 0).cuda()
print(tensor_batch.shape)
output = torch.squeeze(torch.sigmoid(model(tensor_batch)))
np_output = output.cpu().detach().numpy()
# with open("./np_output.pickle", 'wb') as handle:
# pickle.dump(np_output, handle)
# with open("./before_sigmod.pickle", 'wb') as handle:
# pickle.dump(model(tensor_batch).cpu().detach().numpy(), handle)
detected_classes = classes_list[np_output > args.th]
print('done\n')
print(detected_classes)
# print('showing image on screen...')
# fig = plt.figure()
# plt.imshow(im)
# plt.axis('off')
# plt.axis('tight')
# # plt.rcParams["axes.titlesize"] = 10
# plt.title("detected classes: {}".format(detected_classes))
# plt.show()
print('done\n')
def main():
# parsing args
args = parser.parse_args()
# setup model
print('creating model...')
model = create_model(args).cuda()
state = torch.load(args.model_path, map_location='cpu')['model']
model.load_state_dict(state, strict=False)
model.eval()
print('done\n')
# setup data loader
print('creating data loader...')
test_loader = create_dataloader(args)
print('done\n')
# actual validation process
print('doing testing...')
infer(model, test_loader, args)
print('done\n')
def main():
# parsing args
args = parser.parse_args()
# setup model
print('creating model...')
model = create_model(args).cuda()
state = torch.load(args.model_path, map_location='cpu')['model']
model.load_state_dict(state, strict=True)
model.eval()
print('done\n')
# setup data loader
print('creating data loader...')
val_loader = create_dataloader(args)
print('done\n')
# actual validation process
print('doing validation...')
prec1_f, prec5_f = validate(model, val_loader)
print("final top-1 validation accuracy: {:.2f}".format(prec1_f.avg))
def main():
args = parser.parse_args()
args.batch_size = args.batch_size
# setup logger
logger = setup_logger(output=args.output, color=False, name="nuswide")
logger.info("Command: "+' '.join(sys.argv))
# setup model
print('creating and loading the model...')
state = torch.load(args.model_path, map_location='cpu')
args.num_classes = state['num_classes']
args.do_bottleneck_head = False
model = create_model(args).cuda()
model.load_state_dict(state['model'], strict=True)
model.eval()
print('done\n')
# Data loading code
normalize = transforms.Normalize(mean=[0, 0, 0],
std=[1, 1, 1])
val_dataset = NusWideAslDataset(
img_dir=args.img_dir,
csv_path=args.csv_path,
split='val',
transform=transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
normalize,
])
)
print("len(val_dataset)): ", len(val_dataset))
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
validate_multi(val_loader, model, args, logger)
def predict(self):
dataset = create_dataset(
self.opt
) # create a dataset given opt.dataset_mode and other options
model = create_model(
self.opt) # create a model given opt.model and other options
model.setup(
self.opt
) # regular setup: load and print networks; create schedulers
img_rec = None
if self.opt.eval:
model.eval()
for i, data in enumerate(dataset):
model.set_input(data) # unpack data from data loader
model.test() # run inference
visuals = model.get_current_visuals() # get image results
im = visuals['fake_B']
img_rec = util.tensor2im(im)
return img_rec
def main():
# parsing args
args = parser.parse_args()
# gpu setup
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# setup model
print('load pretrained model...')
model = create_model(args)
state = torch.load(args.model_path, map_location='cpu')['model']
model.load_state_dict(state, strict=False)
print('set up the model for fine-grained task...')
for param in model.parameters():
param.requires_grad = False # freeze the convnet
num_in_fc = model.num_features
model.head.fc = nn.Linear(num_in_fc, args.transfer_num_classes)
model = model.to(device)
print('done\n')
# setup data loader
print('creating data loader...')
train_loader = create_dataloader(args)
print('done\n')
# training process
print('start training...')
# only train the fc
print('train only fc layer:')
model = train(model, train_loader, args, device, stage='only_fc')
# train the entire network
print('train entire network:')
for param in model.parameters():
param.requires_grad = True
train(model, train_loader, args, device, stage='full_net')
print('finish training\n')
def main():
# parsing args
args = parser.parse_args()
# setup model
print('creating model...')
model = create_model(args).cuda()
from src.models.tresnet_v2.tresnet_v2 import InplacABN_to_ABN
model2 = InplacABN_to_ABN(model)
aaa = torch.jit.script(model2)
state = torch.load(args.model_path, map_location='cpu')['model']
model.load_state_dict(state, strict=False)
model.eval()
print('done\n')
# setup data loader
print('creating data loader...')
val_loader = create_dataloader(args)
print('done\n')
# actual validation process
print('doing validation...')
prec1_f = validate(model, val_loader)
print("final top-1 validation accuracy: {:.2f}".format(prec1_f.avg))
def main():
args = parser.parse_args()
args.batch_size = args.batch_size
args.do_bottleneck_head = False
# setup model
print('creating model...')
model = create_model(args).cuda()
if args.model_path:
state = torch.load(args.model_path, map_location='cpu')
filtered_dict = {
k: v
for k, v in state['model'].items()
if (k in model.state_dict() and 'head.fc' not in k)
}
model.load_state_dict(filtered_dict, strict=False)
# model.load_state_dict(state['model'], strict=False)
print('done\n')
# Data loading code
normalize = transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
instances_path_val = os.path.join(args.data,
'annotations/instances_val2014.json')
instances_path_train = os.path.join(
args.data, 'annotations/instances_train2014.json')
# data_path_val = os.path.join(args.data, 'val2014')
# data_path_train = os.path.join(args.data, 'train2014')
data_path_val = args.data
data_path_train = args.data
val_dataset = CocoDetection(
data_path_val,
instances_path_val,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
# normalize, # no need, toTensor does normalization
]))
train_dataset = CocoDetection(
data_path_train,
instances_path_train,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
CutoutPIL(cutout_factor=0.5),
RandAugment(),
transforms.ToTensor(),
# normalize,
]))
print("len(val_dataset)): ", len(val_dataset))
print("len(train_dataset)): ", len(train_dataset))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=False)
train_multi_label_coco(model,
train_loader,
val_loader,
args.lr,
gamma_neg=4,
gamma_pos=0,
clip=0.05)
def main():
# parsing args
args = parser.parse_args()
transforms = torchvision.transforms.Compose([
torchvision.transforms.ToPILImage(),
torchvision.transforms.Resize((112, 112)),
torchvision.transforms.ToTensor(),
])
def load_video(video_path: str):
image_size = 112
video = torchvision.io.read_video(video_path,
pts_unit="sec")[0].float()
# [t, n, h, w]
video = video.permute(0, 3, 1, 2)
t_size = 16
#new_video = torch.zeros((min(video.shape[0], t_size), video.shape[1], image_size, image_size))
new_video = torch.zeros(
(t_size, video.shape[1], image_size, image_size))
for t in range(min(video.shape[0], t_size)):
if video.shape[0] < t_size:
frame = transforms(video[t])
new_video[t] = (frame - 127.5) / 127.5
else:
t_ = int(video.shape[0] / t_size * t)
frame = transforms(video[t_])
new_video[t] = (frame - 127.5) / 127.5
return new_video
def validate(model, val_loader):
device = 'cuda'
model.eval()
total, correct = 0, 0
for j, (video, label) in enumerate(val_loader):
#batch = torch.tensor(batch)
#video, label = batch[:,0], batch[:,1]
video, label = video.view(-1, 3, 112, 112), label.view(-1)
video = video.to(device)
#label = label.cuda()
with torch.no_grad():
pred = model(video)
pred = F.softmax(pred, dim=1)
pred = torch.argmax(pred, dim=1).cpu().numpy()
correct += accuracy_score(label, pred, normalize=False)
total += 16 #batch_size
return correct / total * 100
def train(model, train_loader, val_loader, loss_fn, optimizer, scheduler,
max_epochs):
device = 'cuda'
model.train()
loss_fn = loss_fn.to(device)
optimizer.zero_grad()
#loss = torch.as_tensor(0.0).cuda()
best_acc = 0
best_model = copy.deepcopy(model)
for i in range(max_epochs):
for j, (video, label) in enumerate(train_loader):
#batch = torch.tensor(batch)
#video, label = batch[:,0], batch[:,1]
video, label = video.view(-1, 3, 112, 112), label.view(-1)
video = video.to(device)
label = label.to(device)
pred = model(video)
pred = F.softmax(pred, dim=1)
loss = loss_fn(pred, label)
#torch.cuda.empty_cache()
#loss /= len(batch)
loss.backward()
optimizer.step()
scheduler.step()
if j % 100 == 0:
print(f"epoch:{i}, {j}, loss:{loss.item()}")
wandb.log({"loss": loss.item()})
if j % 1000 == 0:
acc = validate(model, val_loader)
wandb.log({"acc": acc})
if acc > best_acc:
best_acc = acc
best_model = copy.deepcopy(model)
print("best_acc:", best_acc)
return best_model
wandb.init(project="stam", name=f"stam-{datetime.datetime.now()}")
# k times cross validation
for i in range(1, 4):
cfg = Config()
device = "cuda"
train_dataset = UCF101(
"./dataset/UCF101",
[
f"./dataset/ucfTrainTestlist/trainlist0{i}.txt",
f"./dataset/ucfTrainTestlist/trainlist0{((i+1)%3)+1}.txt",
],
)
testset = UCF101(
"./dataset/UCF101",
[f"./dataset/ucfTrainTestlist/trainlist0{((i+2)%3)+1}.txt"],
)
t = len(testset)
t, v = int(t * 0.9), t - int(t * 0.9)
test_dataset, val_dataset = torch.utils.data.random_split(
testset, [t, v])
train_loader = DataLoader(
train_dataset,
1,
shuffle=True,
#collate_fn=video_collate,
)
val_loader = DataLoader(
val_dataset,
1,
shuffle=True,
#collate_fn=video_collate,
)
test_loader = DataLoader(
test_dataset,
1,
shuffle=False,
#collate_fn=video_collate,
)
# setup model
print('creating model...')
model = create_model(args).to(device)
#state = torch.load(args.model_path, map_location='cpu')['model']
#model.load_state_dict(state, strict=False)
#model.eval()
print('done\n')
optimizer = torch.optim.Adam(model.parameters(), lr=cfg.base_lr)
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer,
lambda x: (512**-0.5) * min((x + 1)**(-0.5),
(x + 1) * cfg.warmup_steps**(-1.5)),
)
loss_fn = nn.CrossEntropyLoss()
print(sum(p.numel() for p in model.parameters()))
model = train(model,
train_loader,
val_loader,
loss_fn,
optimizer,
scheduler,
max_epochs=cfg.epochs)
torch.save(model, f"./checkpoints/ckpt-{i}.pt")
acc = validate(model, test_loader)
print(f"{i}acc:{acc}")
wandb.log({"test_acc": acc})
def main():
args = parser.parse_args()
args.batch_size = args.batch_size
# setup model
print('creating model...')
#state = torch.load(args.model_path, map_location='cpu')
#args.num_classes = state['num_classes']
args.do_bottleneck_head = True
model = create_model(args).cuda()
ema = EMA(model, 0.999)
ema.register()
#model.load_state_dict(state['model'], strict=True)
#model.train()
classes_list = np.array(list(idx_to_class.values()))
print('done\n')
# Data loading code
normalize = transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
instances_path_val = os.path.join(args.data,
'annotations/instances_val2017.json')
#instances_path_train = os.path.join(args.data, 'annotations/instances_val2017.json')#temprarily use val as train
instances_path_train = os.path.join(
args.data, 'annotations/instances_train2017.json')
data_path_val = os.path.join(args.data, 'val2017')
#data_path_train = os.path.join(args.data, 'val2017')#temporarily use val as train
data_path_train = os.path.join(args.data, 'train2017')
val_dataset = CocoDetection(
data_path_val, instances_path_val,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
normalize,
]))
train_dataset = CocoDetection(
data_path_train, instances_path_train,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
normalize,
]))
print("len(val_dataset)): ", len(val_dataset))
print("len(train_dataset)): ", len(train_dataset))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=False)
criterion = AsymmetricLoss()
params = model.parameters()
optimizer = torch.optim.Adam(params, lr=0.0002,
weight_decay=0.0001) #尝试新的optimizer
total_step = len(train_loader)
scheduler = lr_scheduler.OneCycleLR(optimizer,
max_lr=0.0002,
total_steps=total_step,
epochs=25)
#total_step = len(train_loader)
highest_mAP = 0
trainInfoList = []
Sig = torch.nn.Sigmoid()
#f=open('info_train.txt', 'a')
for epoch in range(5):
for i, (inputData, target) in enumerate(train_loader):
f = open('info_train.txt', 'a')
#model.train()
inputData = inputData.cuda()
target = target.cuda()
target = target.max(dim=1)[0]
#Sig = torch.nn.Sigmoid()
output = Sig(model(inputData))
#output[output<args.thre] = 0
#output[output>=args.thre]=1
#print(output.shape) #(batchsize, channel, imhsize, imgsize)
#print(inputData.shape) #(batchsize, numclasses)
#print(output[0])
#print(target[0])
loss = criterion(output, target)
model.zero_grad()
loss.backward()
optimizer.step()
ema.update()
#store information
if i % 10 == 0:
trainInfoList.append([epoch, i, loss.item()])
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
epoch, 5, i, total_step, loss.item()))
f.write('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}\n'.format(
epoch, 5, i, total_step, loss.item()))
if (i + 1) % 400 == 0:
#储存相应迭代模型
torch.save(
model.state_dict(),
os.path.join('models/',
'model-{}-{}.ckpt'.format(epoch + 1, i + 1)))
#modelName = 'models/' + 'decoder-{}-{}.ckpt'.format(epoch+1, i+1)
mAP_score = validate_multi(val_loader, model, args, ema)
#model.train()
if mAP_score > highest_mAP:
highest_mAP = mAP_score
print('current highest_mAP = ', highest_mAP)
f.write('current highest_mAP = {}\n'.format(highest_mAP))
torch.save(model.state_dict(),
os.path.join('models/', 'model-highest.ckpt'))
f.close()
scheduler.step() #修改学习率
parser.add_argument('--model_path')
parser.add_argument('--model_name', type=str, default='tresnet_m')
parser.add_argument('--num_classes', type=int, default=1000)
parser.add_argument('--input_size', type=int, default=224)
parser.add_argument('--val_zoom_factor', type=int, default=0.875)
parser.add_argument('--batch_size', type=int, default=48)
parser.add_argument('--num_workers', type=int, default=8)
# parsing args
args = parser.parse_args()
val_tfms = create_val_tfm(args)
#### TResNet-M ####
args.model_name = 'tresnet_m'
model_path = './tresnet_m.pth'
model = create_model(args)
state = torch.load(model_path, map_location='cpu')['model']
model.load_state_dict(state, strict=True)
model = InplacABN_to_ABN(model)
model = fuse_bn_recursively(model)
model = model.cuda()
model = model.half()
model.eval()
# Run the benchmark
print('Benchmarking TResNet-M')
for i in range(1): # Two times for caching
ImageNet.benchmark(
model=model,
paper_model_name='TResNet-M-FP16',
paper_arxiv_id='2003.13630',
def validate(args):
# might as well try to validate something
args.prefetcher = not args.no_prefetcher
# create model
model = create_model(
args.model,
num_classes=args.num_classes,
in_chans=3,
global_pool=args.gp)
if args.checkpoint:
load_checkpoint(model, args.checkpoint)
param_count = sum([m.numel() for m in model.parameters()])
_logger.info('Model %s created, param count: %d' % (args.model, param_count))
data_config = resolve_data_config(vars(args), model=model)
# model, test_time_pool = apply_test_time_pool(model, data_config, args)
if torch.cuda.is_available():
model.cuda()
criterion = nn.CrossEntropyLoss().cuda()
dataset = Dataset(args.data)
crop_pct = data_config['crop_pct']
loader = create_loader(
dataset,
input_size=data_config['input_size'],
batch_size=args.batch_size,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
crop_pct=crop_pct)
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
f1_m = AverageMeter()
end = time.time()
total_pred_idx = []
total_truth_idx = []
mistake_image = []
mistake_image_dict = {'calling': [], 'normal': [], 'smoking': [], 'smoking_calling': []}
# class_2_index = {0: 'normal', 1: 'phone', 2: 'smoke'}
class_2_index = {0: 'calling', 1: 'normal', 2: 'smoking', 3: 'smoking_calling'}
with open("./txts/%s.json" % json_name, 'r', encoding="utf-8") as f:
shape_dict = json.load(f)
dets_info = {}
model.eval()
with torch.no_grad():
# warmup, reduce variability of first batch time, especially for comparing torchscript vs non
input = torch.randn((args.batch_size,) + data_config['input_size'])
if torch.cuda.is_available():
input = input.cuda()
model(input)
end = time.time()
for batch_idx, (input, target) in enumerate(loader):
if args.no_prefetcher and torch.cuda.is_available():
target = target.cuda()
input = input.cuda()
# compute output
# t0 = time.time()
output = model(input)
# print("time0: %.8f s" % ((time.time() - t0)))
# t1 = time.time()
# out = output.detach().cpu()
# print("time1: %.8f s" % ((time.time() - t1) / 64))
# print("time2: %.8f s" % ((time.time() - t0) / 64))
# t2 = time.time()
# out = out.cuda().cpu()
# print("time3: %.8f s" % ((time.time() - t2) / 64))
# get prediction index and ground turth index
prob = torch.max(F.softmax(output, -1), -1)[0]
idx = torch.max(F.softmax(output, -1), -1)[1]
target_idx = target.cpu().numpy()
predict_idx = idx.cpu().numpy()
for j in range(len(target_idx)):
total_truth_idx.append(target_idx[j])
total_pred_idx.append(predict_idx[j])
class_dict = loader.dataset.class_to_idx
target_class = list(class_dict.keys())[list(class_dict.values()).index(int(target_idx[j]))]
pred_class = list(class_dict.keys())[list(class_dict.values()).index(int(predict_idx[j]))]
filename = loader.dataset.filenames()[batch_idx * args.batch_size + j]
name = filename.split('/')[-1].split('.')[0]
dets_info[name] = [pred_class, float(prob[j]), shape_dict[name][1], shape_dict[name][2]]
if target_idx[j] != predict_idx[j]:
mistake_image.append(
[loader.dataset.filenames()[batch_idx * args.batch_size + j], target_class, pred_class,
np.round(prob[j].cpu().numpy(), 4)])
mistake_image_dict[class_2_index[predict_idx[j]]].append(
loader.dataset.filenames()[batch_idx * args.batch_size + j])
loss = criterion(output, target)
# measure accuracy and record loss
prec1, prec5 = accuracy(output.data, target, topk=(1, 3))
losses.update(loss.item(), input.size(0))
top1.update(prec1.item(), input.size(0))
top5.update(prec5.item(), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % args.log_freq == 0:
_logger.info(
'Test: [{0:>4d}/{1}] '
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f}) '
'Acc@1: {top1.val:>7.2f} ({top1.avg:>7.2f}) '
'Acc@5: {top5.val:>7.3f} ({top5.avg:>7.3f})'.format(
batch_idx, len(loader), batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg,
loss=losses, top1=top1, top5=top5))
with open("%s/%s.json" % (output_path, json_name.split('-')[0]), "w", encoding="utf-8") as f:
json.dump(dets_info, f)
top1a, top5a = top1.avg, top5.avg
results = OrderedDict(
top1=round(top1a, 4), top1_err=round(100 - top1a, 4),
top5=round(top5a, 4), top5_err=round(100 - top5a, 4),
param_count=round(param_count / 1e6, 2),
img_size=data_config['input_size'][-1],
cropt_pct=crop_pct,
interpolation=data_config['interpolation'],
mistake_image_dict=mistake_image_dict,
pred_idx=total_pred_idx, truth_idx=total_truth_idx)
_logger.info(' * Acc@1 {:.2f} ({:.2f}) Acc@5 {:.2f} ({:.2f})'.format(
results['top1'], results['top1_err'], results['top5'], results['top5_err']))
map, each_ap = eval_map(detFolder="%s/%s.json" % (output_path, json_name.split('-')[0]),
gtFolder="txts/%s.json" % json_name, return_each_ap=True)
_logger.info('Valid mAP: {}, each ap: {}'.format(round(map, 4), each_ap))
return results
def model_train():
"""
Trains model which is used as a feature extractor
:return:None
"""
# Download data
downloadData(data_path='/input/speech_commands/')
# Get data dictionary
dataDict = getDataDict(data_path='/input/speech_commands/')
# Obtain dataframe for each dataset
trainDF = getDataframe(dataDict['train'])
valDF = getDataframe(dataDict['val'])
devDF = getDataframe(dataDict['dev'])
testDF = getDataframe(dataDict['test'])
print("Dataset statistics")
print("Train files: {}".format(trainDF.shape[0]))
print("Validation files: {}".format(valDF.shape[0]))
print("Dev test files: {}".format(devDF.shape[0]))
print("Test files: {}".format(testDF.shape[0]))
# Use TF Data API for efficient data input
train_data, train_steps = getDataset(df=trainDF,
batch_size=BATCH_SIZE,
cache_file='train_cache',
shuffle=True)
val_data, val_steps = getDataset(df=valDF,
batch_size=BATCH_SIZE,
cache_file='val_cache',
shuffle=False)
model = create_model()
model.summary()
# Stop training if the validation accuracy doesn't improve
earlyStopping = EarlyStopping(monitor='val_loss',
patience=PATIENCE,
verbose=1)
# Reduce LR on validation loss plateau
reduceLR = ReduceLROnPlateau(monitor='val_loss',
patience=PATIENCE,
verbose=1)
# Compile the model
model.compile(loss='sparse_categorical_crossentropy',
optimizer=Adam(learning_rate=LEARNING_RATE),
metrics=["sparse_categorical_accuracy"])
# Train the model
history = model.fit(train_data.repeat(),
steps_per_epoch=train_steps,
validation_data=val_data.repeat(),
validation_steps=val_steps,
epochs=EPOCHS,
callbacks=[earlyStopping, reduceLR])
# Save model
print("Saving model")
model.save('../models/marvin_kws.h5')
# Save history data
print("Saving training history")
with open('../models/marvin_kws_history.pickle', "wb") as file:
pickle.dump(history.history, file, protocol=pickle.HIGHEST_PROTOCOL)
plot_history(history=history)
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
torch.manual_seed(args.seed)
model = create_model(args.model,
pretrained=True,
num_classes=args.num_classes,
drop_rate=args.drop,
drop_path_rate=args.drop_path,
drop_block_rate=args.drop_block,
checkpoint_path=args.initial_checkpoint)
if args.local_rank == 0:
_logger.info('Model %s created, param count: %d' %
(args.model, sum([m.numel()
for m in model.parameters()])))
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
if args.num_gpu > 1:
model = nn.DataParallel(model,
device_ids=list(range(args.num_gpu))).cuda()
else:
model.cuda()
optimizer = create_optimizer(args, model)
loss_scaler = None
# optionally resume from a checkpoint
resume_epoch = None
if args.resume:
resume_epoch = resume_checkpoint(
model,
args.resume,
optimizer=None if args.no_resume_opt else optimizer,
loss_scaler=None if args.no_resume_opt else loss_scaler,
log_info=args.local_rank == 0)
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
_logger.info('Scheduled epochs: {}'.format(num_epochs))
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
_logger.error(
'Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(mixup_alpha=args.mixup,
cutmix_alpha=args.cutmix,
cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob,
switch_prob=args.mixup_switch_prob,
elementwise=args.mixup_elem,
label_smoothing=args.smoothing,
num_classes=args.num_classes)
if args.prefetcher:
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_workers=args.workers,
collate_fn=collate_fn,
)
eval_dir = os.path.join(args.data, 'valid')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
_logger.error(
'Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
num_workers=args.workers,
crop_pct=data_config['crop_pct'],
)
if mixup_active:
# smoothing is handled with mixup target transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing).cuda()
train_loss_ce = nn.CrossEntropyLoss().cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
plateau_num = 0
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(model=model,
optimizer=optimizer,
args=args,
amp_scaler=loss_scaler,
checkpoint_dir=output_dir,
recovery_dir=output_dir,
decreasing=decreasing,
max_history=2)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
with open("./txts/%s.json" % json_name, 'r', encoding="utf-8") as f:
shape_dict = json.load(f)
try:
for epoch in range(start_epoch, num_epochs):
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
[train_loss_fn, train_loss_ce],
args,
lr_scheduler=lr_scheduler,
output_dir=output_dir,
mixup_fn=mixup_fn)
eval_metrics, dets_info = validate(model,
loader_eval,
validate_loss_fn,
args,
shape_dict=shape_dict)
with open("%s/v.json" % output_dir, "w", encoding="utf-8") as f:
json.dump(dets_info, f)
map = round(
eval_map(detFolder="%s/v.json" % output_dir,
gtFolder="txts/%s.json" % json_name), 4)
eval_metrics["map"] = map
_logger.info('Valid mAP: {}'.format(map))
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
saver.save_prefix = "%.2f-%s" % (eval_metrics["top1"], map)
best_metric, best_epoch = saver.save_checkpoint(
epoch, metric=save_metric)
if eval_metrics[eval_metric] >= best_metric:
plateau_num = 0
else:
plateau_num += 1
# 超过30个epoch还没有更新metric，停止运行
if plateau_num == 30:
break
except KeyboardInterrupt:
pass
if best_metric is not None:
_logger.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def create_ensembles(shape, l_name, contamination=0.01):
"""
Utility function for creating the ensembles
:param shape: The input shape
:param l_name: The idenfitier of the local outlier detector
:param contamination: The contamination (for some models this is a parameter)
:return: array(C), array(L)
"""
num_clients = shape[0]
c = create_models(num_clients, shape[-1], compression_factor=0.4)
l = None
if l_name == "lof1":
l = [
LocalOutlierFactor(n_neighbors=1,
contamination=contamination,
novelty=True) for _ in range(num_clients)
]
if l_name == "lof2":
l = [
LocalOutlierFactor(n_neighbors=2,
contamination=contamination,
novelty=True) for _ in range(num_clients)
]
if l_name == "lof4":
l = [
LocalOutlierFactor(n_neighbors=4,
contamination=contamination,
novelty=True) for _ in range(num_clients)
]
if l_name == "lof8":
l = [
LocalOutlierFactor(n_neighbors=8,
contamination=contamination,
novelty=True) for _ in range(num_clients)
]
if l_name == "lof16":
l = [
LocalOutlierFactor(n_neighbors=16,
contamination=contamination,
novelty=True) for _ in range(num_clients)
]
if l_name == "lof32":
l = [
LocalOutlierFactor(n_neighbors=32,
contamination=contamination,
novelty=True) for _ in range(num_clients)
]
if l_name == "lof64":
l = [
LocalOutlierFactor(n_neighbors=64,
contamination=contamination,
novelty=True) for _ in range(num_clients)
]
if l_name == "lof100":
l = [
LocalOutlierFactor(n_neighbors=100,
contamination=contamination,
novelty=True) for _ in range(num_clients)
]
if l_name == "xstream":
l = [Chains(k=100, nchains=100, depth=15) for _ in range(num_clients)]
if l_name == "ae":
l = [
create_model(shape[-1], compression_factor=0.4)
for _ in range(num_clients)
]
if l_name == "if":
l = [
IsolationForest(contamination=contamination)
for _ in range(num_clients)
]
if not l:
raise KeyError("No valid local outlier detector name provided.")
return np.array(c), np.array(l)
else:
data_path = "/home/data/classification/action/new_data/test/"
if data_path.split('/')[-1] == 'valid':
dataset = Dataset(data_path,
transform=tta_test_transform(img_size))
save_weights = True
else:
dataset = TestDataset(data_path,
transform=tta_test_transform(img_size))
save_weights = False
checkpoint = glob.glob(
os.path.join(output_root,
checkpoint_list[index] + '/*best*.pth.tar'))[0]
model = create_model('%s' % model_name.split('-')[-2],
num_classes=args.num_classes,
in_chans=3,
checkpoint_path='%s' % checkpoint)
model = model.cuda()
model.eval()
test_loader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
model.default_cfg['model_name'] = model_list[index]
prediction_output, true_output = perform_predict(
model,
test_loader,
model_weight,
def main():
args = parser.parse_args()
args.do_bottleneck_head = False
# setup dist training
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
print('Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.' % (args.rank, args.world_size))
else:
print('Training with a single process on 1 GPUs.')
assert args.rank >= 0
# setup logger
logger = setup_logger(output=args.output, distributed_rank=dist.get_rank(), color=False, name="Coco")
logger.info("Command: "+' '.join(sys.argv))
if dist.get_rank() == 0:
path = os.path.join(args.output, "config.json")
with open(path, 'w') as f:
json.dump(vars(args), f, indent=2)
logger.info("Full config saved to {}".format(path))
os.makedirs(osp.join(args.output, 'tmpdata'), exist_ok=True)
# Setup model
logger.info('creating model...')
model = create_model(args).cuda()
if args.model_path: # make sure to load pretrained ImageNet model
state = torch.load(args.model_path, map_location='cpu')
filtered_dict = {k: v for k, v in state['model'].items() if
(k in model.state_dict() and 'head.fc' not in k)}
model.load_state_dict(filtered_dict, strict=False)
logger.info('done\n')
ema = ModelEma(model, 0.9997) # 0.9997^641=0.82
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], broadcast_buffers=False)
# COCO Data loading
instances_path_val = os.path.join(args.data, 'annotations/instances_val2014.json')
instances_path_train = os.path.join(args.data, 'annotations/instances_train2014.json')
data_path_val = args.data
data_path_train = args.data
val_dataset = CocoDetection(data_path_val,
instances_path_val,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
# normalize, # no need, toTensor does normalization
]))
train_dataset = CocoDetection(data_path_train,
instances_path_train,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
CutoutPIL(cutout_factor=0.5),
RandAugment(),
transforms.ToTensor(),
# normalize,
]))
logger.info("len(val_dataset)): {}".format(len(val_dataset)))
logger.info("len(train_dataset)): {}".format(len(train_dataset)))
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False)
assert args.batch_size // dist.get_world_size() == args.batch_size / dist.get_world_size(), 'Batch size is not divisible by num of gpus.'
# Pytorch Data loader
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size // dist.get_world_size(),
shuffle=not args.distributed,
num_workers=args.workers, pin_memory=True, sampler=train_sampler, drop_last=True)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size // dist.get_world_size(),
shuffle=False,
num_workers=args.workers, pin_memory=False, sampler=val_sampler)
# Actuall Training
train_multi_label_coco(model, ema, train_loader, val_loader, args.lr, args, logger)
def main():
setup_default_logging()
args = parser.parse_args()
# might as well try to do something useful...
args.pretrained = args.pretrained or not args.checkpoint
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
args.checkpoint = glob.glob(args.checkpoint + '/*.pth')[0]
# create model
model = create_model(args.model,
num_classes=args.num_classes,
in_chans=3,
pretrained=args.pretrained)
load_checkpoint(model, args.checkpoint)
logging.info('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
args.img_size = int(args.checkpoint.split('/')[-2].split('-')[-1])
config = resolve_data_config(vars(args), model=model)
# model, test_time_pool = apply_test_time_pool(model, config, args)
if torch.cuda.is_available():
model = model.cuda()
loader = create_loader(Dataset(args.data),
input_size=config['input_size'],
batch_size=args.batch_size,
use_prefetcher=False,
interpolation=config['interpolation'],
mean=config['mean'],
std=config['std'],
num_workers=args.workers,
crop_pct=config['crop_pct'])
model.eval()
batch_time = AverageMeter()
end = time.time()
topk_ids = []
scores = []
total_pred_idx = []
total_truth_idx = []
with torch.no_grad():
for batch_idx, (input, target) in enumerate(loader):
if torch.cuda.is_available():
input = input.cuda()
output = model(input)
prob = torch.max(F.softmax(output, -1), -1)[0]
idx = torch.max(F.softmax(output, -1), -1)[1]
total_pred_idx.extend(idx.cpu().numpy())
total_truth_idx.extend(target.cpu().numpy())
scores.extend(prob.cpu().numpy())
topk_ids.extend(idx.cpu().numpy())
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % args.log_freq == 0:
logging.info(
'Predict: [{0}/{1}] Time {batch_time.val:.3f} ({batch_time.avg:.3f})'
.format(batch_idx, len(loader), batch_time=batch_time))
# result_file_path = os.path.join('./results', model_name)
# if not os.path.exists(result_file_path):
# os.makedirs(result_file_path)
# res_cf = open('%s/results-all.csv' % result_file_path, mode='w')
# for i in range(len(total_pred_idx)):
# res_cf.write('{0},'.format(str(total_pred_idx[i])))
# res_cf.write('\n')
# for i in range(len(total_truth_idx)):
# res_cf.write('{0},'.format(str(total_truth_idx[i])))
# dst_root = './infer/%s' % args.checkpoint.split('/')[-2]
# if not os.path.exists(dst_root):
# os.makedirs(dst_root)
# else:
# shutil.rmtree(dst_root)
result_list = []
# class_2_index = {0: 'normal', 1: 'calling', 2: 'smoking'}
class_2_index = {
0: 'calling',
1: 'normal',
2: 'smoking',
3: 'smoking_calling'
}
with open(os.path.join(args.output_dir, 'result.json'),
'w',
encoding="utf-8") as out_file:
filenames = loader.dataset.filenames()
for i in range(len(scores)):
filename = filenames[i].split('/')[-1]
name = class_2_index[topk_ids[i]]
result_data = {
"image_name": str(filename),
"category": name,
"score": scores[i]
}
result_list.append(result_data)
# if scores[i] > 0.95:
# dst_path = os.path.join(dst_root, name)
# if not os.path.exists(dst_path):
# os.makedirs(dst_path)
# shutil.copy(filenames[i], os.path.join(dst_path, filename))
json.dump(result_list, out_file, cls=MyEncoder, indent=4)
from src.data import create_dataset
from src.models import create_model
from src.util.visualizer import save_images_to_path
import sys
if __name__ == '__main__':
opt = TestOptions().parse() # get test options
# hard-code some parameters for test
opt.num_threads = 0 # test code only supports num_threads = 1
opt.batch_size = 1 # test code only supports batch_size = 1
opt.serial_batches = True # disable data shuffling; comment this line if results on randomly chosen images are needed.
opt.no_flip = True # no flip; comment this line if results on flipped images are needed.
opt.display_id = -1 # no visdom display; the test code saves the results to a HTML file.
dataset = create_dataset(
opt) # create a dataset given opt.dataset_mode and other options
model = create_model(
opt) # create a model given opt.model and other options
model.setup(
opt) # regular setup: load and print networks; create schedulers
if opt.force_test_output == '':
print('--force_test_output must be defined.')
sys.exit()
if opt.eval:
model.eval()
for i, data in enumerate(dataset):
if i >= opt.num_test: # only apply our model to opt.num_test images.
break
model.set_input(data) # unpack data from data loader
model.test() # run inference
visuals = model.get_current_visuals() # get image results
def main():
args = parser.parse_args()
args.batch_size = args.batch_size
# setup model
print('creating model...')
state = torch.load(args.model_path, map_location='cuda')
#args.num_classes = state['num_classes']
args.do_bottleneck_head = False
model = create_model(args).cuda()
#load pretrained model
model.load_state_dict(state['model'], strict=True)
in_features = model.head.fc.in_features
#initially, num_class == 1000 in order to load model pretrained on ImageNet correctly, not change it to 80
model.head.fc = torch.nn.Linear(in_features, 80).cuda()
ema = EMA(model, 0.999)
ema.register()
model.train()
classes_list = np.array(list(idx_to_class.values()))
print('done\n')
# Data loading code
normalize = transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
instances_path_val = os.path.join(args.data,
'annotations/instances_val2017.json')
#instances_path_train = os.path.join(args.data, 'annotations/instances_val2017.json')#temprarily use val as train
instances_path_train = os.path.join(
args.data, 'annotations/instances_train2017.json')
data_path_val = os.path.join(args.data, 'val2017')
#data_path_train = os.path.join(args.data, 'val2017')#temporarily use val as train
data_path_train = os.path.join(args.data, 'train2017')
val_dataset = CocoDetection(
data_path_val,
instances_path_val,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
#RandAugment(),
transforms.ToTensor(),
normalize,
]))
train_dataset = CocoDetection(
data_path_train,
instances_path_train,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
RandAugment(),
transforms.ToTensor(),
normalize,
#Cutout(n_holes = 1, length = 16)
Cutout(n_holes=1, length=args.image_size / 2)
]))
print("len(val_dataset)): ", len(val_dataset))
print("len(train_dataset)): ", len(train_dataset))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=False)
#set optimizer
lr = 0.0002
Epoch = 10
weight_decay = 0.0001
criterion = AsymmetricLoss()
#params = model.parameters()
skip = {}
if hasattr(model, 'no_weight_decay'):
skip = model.no_weight_decay()
parameters = add_weight_decay(model, weight_decay, skip)
weight_decay = 0.
opt_args = dict(lr=lr, weight_decay=weight_decay)
#optimizer = torch.optim.Adam(params, lr=lr)#尝试新的optimizer
optimizer = AdamW(parameters, **opt_args)
total_step = len(train_loader)
scheduler = lr_scheduler.OneCycleLR(optimizer,
max_lr=lr,
total_steps=total_step,
epochs=Epoch)
#total_step = len(train_loader)
highest_mAP = 0
trainInfoList = []
Sig = torch.nn.Sigmoid()
#f=open('info_train.txt', 'a')
for epoch in range(Epoch):
for i, (inputData, target) in enumerate(train_loader):
f = open('info_train.txt', 'a')
#model.train()
inputData = inputData.cuda()
target = target.cuda()
target = target.max(dim=1)[0]
#Sig = torch.nn.Sigmoid()
output = Sig(model(inputData))
#output[output<args.thre] = 0
#output[output>=args.thre]=1
#print(output.shape) #(batchsize, channel, imhsize, imgsize)
#print(inputData.shape) #(batchsize, numclasses)
#print(output[0])
#print(target[0])
loss = criterion(output, target)
model.zero_grad()
loss.backward()
optimizer.step()
ema.update()
#store information
if i % 10 == 0:
trainInfoList.append([epoch, i, loss.item()])
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
epoch, Epoch, i, total_step, loss.item()))
f.write('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}\n'.format(
epoch, Epoch, i, total_step, loss.item()))
if (i + 1) % 100 == 0:
#储存相应迭代模型
torch.save(
model.state_dict(),
os.path.join('models/',
'model-{}-{}.ckpt'.format(epoch + 1, i + 1)))
f.close()
scheduler.step() #修改学习率
#validate per epoch and record result
f = open('info_train.txt', 'a')
model.eval()
mAP_score = validate_multi(val_loader, model, args, ema)
model.train()
if mAP_score > highest_mAP:
highest_mAP = mAP_score
print('current highest_mAP = ', highest_mAP)
f.write('current_mAP = {}, highest_mAP = {}\n'.format(
mAP_score, highest_mAP))
torch.save(model.state_dict(),
os.path.join('models/', 'model-highest.ckpt'))
f.close()
