def main(args):
print('----------------------------------------------------')
print("{}-way-{}-shot Few-Shot Relation Classification".format(
args.N, args.K))
print("Model: {}".format(args.Model))
print("config:", args)
print('----------------------------------------------------')
start_time = time.time()
mymodel = MyModel(args)
mymodel_clone = MyModel_Clone(args)
best_acc = 0.0
best_loss = 0.0
for file_name in os.listdir('model_checkpoint'):
if 'isNPM.tar' in file_name:
model_file = 'model_checkpoint/' + file_name
mymodel.load_state_dict(torch.load(model_file))
acc, loss = test_model(mymodel, mymodel_clone, args)
print('model_name:', model_file)
print('[TEST] | loss: {0:2.6f}, accuracy: {1:2.2f}%'.format(
loss, acc * 100))
if acc > best_acc:
best_acc = acc
best_loss = loss
best_model_file = model_file
print('best_model_name:', best_model_file)
print('best_loss:', best_loss)
print('best_acc:', best_acc)
def main():
device = torch.device('cpu')
torch.manual_seed(1234)
layer = 2
size = 1024
func = "relu"
model = MyModel(layer, size)
model.load_state_dict(torch.load('model_state_dict_final'))
model.eval()
transforms = T.Compose([T.ToTensor(), T.Normalize((0.5, ), (0.5, ))])
test_dataset = MnistDatasetTest('data', 'test', transforms)
test_dataloader = DataLoader(test_dataset,
batch_size=64,
shuffle=False,
num_workers=4)
result = open("result.txt", "w") # to create a result
with torch.no_grad():
for images, image_name in test_dataloader:
images = images.to(device)
prediction = model(images, layer, func)
for i in range(images.size()[0]):
x = image_name[i] + ' ' + str(int(torch.argmax(prediction[i])))
result.write(x)
result.write("\n")
result.close()
def _load(self, filePath):
checkpoint = torch.load(filePath)
model = MyModel(
device, checkpoint['inputSize'], checkpoint['gatedCnnOutputSize'],
checkpoint['gatedCnnStride1'], checkpoint['gatedCnnStride2'],
checkpoint['gatedCnnKernel1'], checkpoint['gatedCnnKernel2'],
checkpoint['lstmLayer'], checkpoint['lstmHiddenSize'],
checkpoint['fcOutputSize'], checkpoint['dropout'])
model.load_state_dict(checkpoint['stateDict'])
model.eval()
if self.device.type == 'cpu':
model.cpu()
else:
model.cuda(device=self.device)
return model
NUM_EPISODES = 4000
TEST_INTERVAL = 200
PRINT_INTERVAL = 50
LEARNING_RATE = 1e-5
if USE_MEMORY:
LEARNING_RATE = (4e-4 if DOUBLE_Q else 3e-3)
RENDER_INTERVAL = 20
ENV_NAME = 'CartPole-v0'
env = gym.make(ENV_NAME)
state_shape = len(env.reset())
n_actions = env.action_space.n
model = MyModel(state_shape, n_actions).to(device)
target = MyModel(state_shape, n_actions).to(device)
target.load_state_dict(model.state_dict())
target.eval()
optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)
loss_function = torch.nn.MSELoss(reduction='mean')
memory = ReplayBuffer()
def choose_action(state, test_mode=False):
if random.random() < EPS_EXPLORATION:
return env.action_space.sample()
return model.select_action(torch.from_numpy(state)).item()
# Given a tuple (s_t, a_t, s_{t+1}, r_t, done_t) update your model weights
def optimize_model(state, action, next_state, reward, done):
# For double Q, use the other Q function for getting values of the next state.
nextStateQ = target if DOUBLE_Q else model
import numpy as np
import torch
import torch.nn as nn
from SMPL.smpl_np import SMPLModel
from model import MyModel
input_size = 14 * 2
output_size = 24 * 3
smpl = SMPLModel('model/feman_lbs.pkl')
model = MyModel(input_size, output_size)
model.double()
model.eval()
model.load_state_dict(
torch.load('trained_model/2018_08_20/epoch_46_model.ckpt'))
joint_t = np.load('/mnt/data/dataset/SMPL/test/joint/joint_0000057.npy')
pose_t = np.load('/mnt/data/dataset/SMPL/test/pose/pose_0000057.npy')
joint_t = torch.from_numpy(joint_t[:, 0:2])
joint_t = joint_t.reshape(1, -1)
pose_p = model(joint_t).data.numpy()
pose_p = pose_p.reshape(24, 3)
smpl.set_params(pose=pose_t)
smpl.save_to_obj('test_result/vertices_t.obj')
smpl.save_joints_to_obj('test_result/joint_t.obj')
smpl.set_params(pose=pose_p)
smpl.save_to_obj('test_result/vertices_p.obj')
smpl.save_joints_to_obj('test_result/joint_p.obj')
# if np.allclose(joint_t, smpl.get_changed_joints()[]):
def main(args):
model = MyModel(args.size)
#model = get_vgg_model()
print(model)
## TODO: Crop method
if args.do_train:
transformer = get_train_transform(512)
dev_transformer = get_predict_transform(512)
start_epoch = 0
if args.load_model:
logging.info('loading model...... %s' % (args.load_model))
checkpoint = torch.load(args.load_model)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch'] + 1
### trainset
trainset = MyDataset(args.data_list,
args.data_dir,
transform=transformer)
train_loader = DataLoader(dataset=trainset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=collate_picture)
### dev set
devset = MyDataset(args.dev_file,
args.data_dir,
transform=dev_transformer)
dev_loader = DataLoader(dataset=devset,
batch_size=args.batch_size,
shuffle=False,
collate_fn=collate_picture)
train(args, model, train_loader, dev_loader, start_epoch)
#### predict
if args.do_predict:
transformer = get_predict_transform(512)
if args.load_model == None and not args.do_train:
logging.error('load model error')
exit(1)
elif args.load_model:
logging.info('loading model...... %s' % (args.load_model))
checkpoint = torch.load(args.load_model)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch'] + 1
#model.load_state_dict(torch.load(args.load_model))
if args.predict_file == None:
logging.error('No predict file')
exit(1)
logging.info('start reading predict files')
predictset = MyDataset(args.predict_file,
args.data_dir,
transform=transformer)
print(predictset[0][0].shape)
print('len', len(predictset))
predict_loader = DataLoader(dataset=predictset,
batch_size=args.batch_size,
shuffle=False,
collate_fn=collate_picture)
predict(args, model, predict_loader)
i, num_test_episodes, episode_total_reward))
break
test_rewards.append(episode_total_reward)
test_env.close()
return sum(test_rewards) / num_test_episodes
if __name__ == "__main__":
import argparse
from model import MyModel
parser = argparse.ArgumentParser()
parser.add_argument('--model-path',
default=None,
type=str,
help='Path to the model weights.')
parser.add_argument('--env',
default=None,
type=str,
help='Name of the environment.')
args = parser.parse_args()
env = gym.make(args.env)
model = MyModel(state_size=len(env.reset()),
action_size=env.action_space.n)
model.load_state_dict(torch.load(args.model_path))
model = model.to(device)
env.close()
eval_policy(policy=model, env=args.env, render=True, verbose=True)
'''
import os
import torch
import torch.nn as nn
from torch.optim import Adam
import torch.nn.functional as F
#from dataset import ws
from lib import ws, max_len
from model import MyModel
from dataset import get_dataloader
model = MyModel()
optimizer = Adam(model.parameters(), 0.001)
if os.path.exists('./model/model1.pkl'):
model.load_state_dict(torch.load('./model/model1.pkl'))
optimizer.load_state_dict(torch.load('./model/optimizer1.pkl'))
def train(epoch):
for idx, (input, target) in enumerate(get_dataloader):
optimizer.zero_grad()
output = model(input)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
print(epoch, idx, loss.item())
if idx % 100 == 0:
torch.save(model.state_dict(), ('./model/model1.pkl'))
torch.save(optimizer.state_dict(), ('./model/optimizer1.pkl'))
def predict(dim,
names,
weight,
batch_size,
pretrain_model_path,
model_types=None):
print('-' * 100)
print('multi-models begin predicting ...')
print('-' * 100)
# read test data
test_file = '../data/Dataset/test.csv'
# data
test_df = pd.read_csv(test_file)
test_ids = test_df['id'].values.tolist()
result_prob_tmp = torch.zeros((len(test_ids), 2))
# load model
for i, name in enumerate(names):
# 3.17 add
weight_ = weight[i]
model_path = '../user_data/model_store/' + name + '.pkl'
state = torch.load(model_path)
# 3.10 add
model_type = model_types[i]
if model_type == 'mlp':
test_iter = MyDataset(file=test_file, is_train=False, pretrain_model_path=pretrain_model_path[i])
test_iter = get_dataloader(test_iter, batch_size, shuffle=False, drop_last=False)
model = MyModel(dim=dim[i], pretrain_model_path=pretrain_model_path[i])
elif model_type == 'cnn':
test_iter = MyDataset(file=test_file, is_train=False, pretrain_model_path=pretrain_model_path[i])
test_iter = get_dataloader(test_iter, batch_size, shuffle=False, drop_last=False)
model = MyTextCNNModel(dim=dim[i], pretrain_model_path=pretrain_model_path[i])
elif model_type == 'rcnn':
test_iter = MyDataset(file=test_file, is_train=False, pretrain_model_path=pretrain_model_path[i])
test_iter = get_dataloader(test_iter, batch_size, shuffle=False, drop_last=False)
model = MyRCNNModel(dim=dim[i], pretrain_model_path=pretrain_model_path[i])
model.to(device)
model.load_state_dict(state['model_state'])
model.eval()
print('-'*20, 'model', i, '-'*20)
print('load model:%s, loss:%.4f, e:%d, lr:%.7f, time:%d' %
(name, state['loss'], state['e'], state['lr'], state['time']))
# predict
with torch.no_grad():
j = 0
for batch in tqdm(test_iter):
batch = [b.cuda() for b in batch]
out = model(batch, task='eval')
out = out.cpu() # gpu -> cpu
if j == 0:
tmp = out # 初始化 tmp
else:
tmp = torch.cat([tmp, out], dim=0) # 将之后的预测结果拼接到 tmp 中
j += 1
# 当前 模型预测完成
print('model', i, 'predict finished!\n')
# 3.17 按权重融合
result_prob_tmp += (weight_ / len(names)) * tmp
# 删除模型
del model
gc.collect()
time.sleep(1)
# 3.10 当前融合策略：prob 简单的取 avg
_, result = torch.max(result_prob_tmp, dim=-1)
result = result.numpy()
# 3.16 update: label 0的prob 大于 3，就认为是 label=0
# with open('tmp.txt', 'w', encoding='utf-8') as f:
# for r in result_prob_tmp:
# f.write(str(r) + '\n')
# save result
df = pd.DataFrame()
df['id'] = test_ids
df['label'] = result
df.to_csv(("../prediction_result/result_"+datetime.datetime.now().strftime('%Y%m%d_%H%M%S') + ".csv"), encoding='utf-8', index=False)
def test(rank, cfg, loader, dataset, epoch, uuids, nls, scene_threshold,
total_threshold):
dataset.load_frame = False
dist.init_process_group(backend="nccl",
rank=rank,
world_size=cfg.num_gpu,
init_method="env://")
torch.cuda.set_device(rank)
cudnn.benchmark = True
model = MyModel(cfg,
len(dataset.nl),
dataset.nl.word_to_idx['<PAD>'],
norm_layer=nn.BatchNorm2d,
num_colors=len(CityFlowNLDataset.colors),
num_types=len(CityFlowNLDataset.vehicle_type) - 2).cuda()
model = DistributedDataParallel(model,
device_ids=[rank],
output_device=rank,
broadcast_buffers=cfg.num_gpu > 1,
find_unused_parameters=False)
saved_dict = torch.load(f'save/{epoch}.pth',
map_location=torch.device(f'cuda:{rank}'))
model.load_state_dict(saved_dict, True)
model.eval()
final_results = {}
a = len(nls) // cfg.num_gpu
start = a * rank
end = None if (rank + 1) == cfg.num_gpu else a * (rank + 1)
end_str = 'end' if end == None else end
print(f'process number: {rank}, {start}:{end_str}')
for nlidx, (uuid,
query_nl) in enumerate(zip(uuids[start:end], nls[start:end])):
nlidx = nlidx + start
print(f'{nlidx} / {len(nls)}')
cache_nl = torch.load(f'cache/{epoch}/{uuid}.pth',
map_location=torch.device(f'cuda:{rank}'))
cache_nl, vehicle_type, vehicle_color = cache_nl['nls'], cache_nl[
'type'], cache_nl['color']
# nls = []
# for nl in query_nl:
# nl = torch.tensor(dataset.nl.sentence_to_index(nl, is_train=False)).cuda()
# nls.append(nl.unsqueeze(0).transpose(1, 0))
uuids_per_nl = []
prob_per_nl = []
for idx, (id, frames, boxes, paths, rois, labels) in enumerate(loader):
# print(f'{nlidx}_{idx}')
with torch.no_grad():
boxes = boxes.squeeze(0).numpy()
rois = rois.squeeze(0).numpy()
# print(rois)
frames = frames.squeeze(0)
# print(frames.shape)
# b = frames.shape[0]
labels = labels.squeeze(0)
labels = labels.cuda()
cache = torch.load(f'cache/{epoch}/{idx}_0.pth',
map_location=torch.device(f'cuda:{rank}'))
# print(cache)
frame, cs, vs = cache
# if vehicle_type != -1 and vs != vehicle_type:
# continue
# if vehicle_color != -1 and cs != vehicle_color:
# continue
# print(frame.device)
# results = []
nl1 = cache_nl[0]
nl2 = cache_nl[1]
nl3 = cache_nl[2]
bs = frame.shape[0]
# cache = cache[:num_of_vehicles]
# print(cache.shape)
if nl1.shape[0] != bs:
nl1 = cache_nl[0].expand(bs, -1, -1).cuda()
nl2 = cache_nl[1].expand(bs, -1, -1).cuda()
nl3 = cache_nl[2].expand(bs, -1, -1).cuda()
am1 = model(nl1, frame, labels)
am2 = model(nl2, frame, labels)
am3 = model(nl3, frame, labels)
# am1, c1, v1 = model(nl1, frame, labels)
# am2, c2, v2 = model(nl2, frame, labels)
# am3, c3, v3 = model(nl3, frame, labels)
activation_aggregation = (am1 + am2 + am3) / 3
# c_aggregation = (c1 + c2 + c3) / 3
# v_aggregation = (v1 + v2 + v3) / 3
# activation_aggregation = model(nl1, frame) +\
# model(nl2, frame) +\
# model(nl3, frame)
# activation_aggregation = activation_aggregation / 3
# results.append(activation_aggregation)
# cs.append(c_aggregation)
# vs.append(v_aggregation)
results = activation_aggregation.cpu().numpy()
# cs = c_aggregation.mean(dim=0).cpu().numpy()
# vs = v_aggregation.mean(dim=0).cpu().numpy()
# version 1
# cache = torch.load(f'cache/{epoch}/{idx}.pth')
# results = []
# for batch_idx in range(cache.shape[0]):
# output = model(cache_nl[0], cache[batch_idx:batch_idx+1]).sigmoid()
# results.append(output.squeeze(0).cpu().detach().numpy())
prob = compute_probability_of_activations(
results, rois, scene_threshold)
# if vehicle_type != -1 and np.argmax(vs) != vehicle_type:
# prob = 0.
# if vehicle_color != -1 and np.argmax(cs) != vehicle_color:
# prob = 0.
###### visualization
# if not os.path.exists('results/' + query_nl[0]):
# os.mkdir('results/' + query_nl[0])
# # cs = np.argmax(cs)
# cs = cs.item()
# vs = vs.item()
# cs = CityFlowNLDataset.colors[cs]
# # vs = np.argmax(vs)
# vs = CityFlowNLDataset.vehicle_type[vs]
# if prob > total_threshold:
# print(f'color: {cs}, type: {vs}')
# save_img(np.squeeze(results[0], axis=0) * 255, cv2.imread(paths[0][0]), boxes[0], f"results/{query_nl[0]}/{idx}_{prob}.png")
###### end visualization
# for submission
uuids_per_nl.append(id[0])
prob_per_nl.append(prob)
final_results['uuids_order'] = uuids_per_nl
final_results[uuid] = prob_per_nl
# uuids_per_nl = np.array(uuids_per_nl)
# # print(uuids_per_nl.shape)
# prob_per_nl = np.array(prob_per_nl)
# prob_per_nl_arg = (-prob_per_nl).argsort(axis=0)
# sorted_uuids_per_nl = uuids_per_nl[prob_per_nl_arg]
# # print(prob_per_nl[prob_per_nl_arg])
# final_results[uuid] = sorted_uuids_per_nl.tolist()
# print(len(final_results.keys()))
with open(f'results/submit_{epoch}_{start}_{end_str}.json', 'w') as fp:
json.dump(final_results, fp)
def infer(PATH_IMG_INFER, PATH_MODEL_PARAMS, PATH_INFER_SAVE, DEPTH,
NUM_CHANNELS, MULT_CHAN, NUM_CLASSES, Z_RANGE, IMG_WIDTH,
IMG_HEIGHT):
### instantiate model and load the model parameters
my_model = MyModel(n_in_channels=NUM_CHANNELS,
mult_chan=MULT_CHAN,
depth=DEPTH)
my_model.load_state_dict(
torch.load(PATH_MODEL_PARAMS, map_location=torch.device('cpu')))
my_model.eval()
### read in the tiff and place into a ZStack object
stack = ZStack(path=PATH_IMG_INFER, zrange=Z_RANGE)
# separate zslices; assign each to its own Slice object, and hold them all in a list
slices = [Slice(x) for x in stack.images]
# split each slice into tiles
tiles = np.array(
[slices[z].tile(IMG_WIDTH, IMG_HEIGHT) for z in range(len(slices))])
# normalize (z-score) the tiles before passing to model
[[tiles[y][x].normalize() for x in range(0, tiles.shape[1])]
for y in range(0, tiles.shape[0])]
# reshape data; h=z_slices, w=numtiles for each z-slice
h, w = np.shape(tiles)
flat_tiles = np.zeros((h * w, 1, IMG_WIDTH, IMG_HEIGHT), dtype=np.float32)
for i in range(h * w):
flat_tiles[i] = tiles.flatten(order='C')[i].reshape(
IMG_WIDTH, IMG_HEIGHT).slice
# load data into pytorch data structs
dataset = FormsDataset(flat_tiles, num_classes=NUM_CLASSES)
hold = []
data_loader = DataLoader(dataset, batch_size=1, shuffle=False)
for i, images in enumerate(data_loader, 1):
images = images.type(torch.FloatTensor)
hold.append(images)
# now lets perform inference on the tiles, store them in 'predictions'
predictions = []
for tile in hold:
predictions.append(my_model(tile))
# and lets time ourselves
print("--- Inference: {} seconds ---".format(time.time() - start_time))
### now we have to stitch our images back together!
unflat_predictions = np.array(predictions).reshape((h, w))
stack_infer = []
for z in unflat_predictions:
stack_infer.append(
stitch(z, IMG_WIDTH=IMG_WIDTH, IMG_HEIGHT=IMG_HEIGHT).toimage())
stack_infer[0].save(PATH_INFER_SAVE,
save_all=True,
append_images=stack_infer[1:])
print("--- Finish: {} seconds ---".format(time.time() - start_time))
class MultiHeadTester:
def __init__(
self,
dataset_path='./drive/MyDrive/datasets/car classification/train_data',
batch_size=1,
model_name='tf_efficientnet_b3_ns',
test_csv='./train_labels.csv',
unique_csv='./train_labels.csv',
output_dir='../drive/MyDrive/ckpt/grapheme/submission.csv',
ckpt='../drive/MyDrive/ckpt/grapheme/20.pth'):
# initialize attributes
self.dataset_path = dataset_path
self.batch_size = batch_size
self.model_name = model_name
self.test_csv = test_csv
self.unique_csv = unique_csv
self.output_dir = output_dir
self.ckpt = ckpt
if model_name == 'tf_efficientnet_b0_ns':
self.input_size = (224, 224)
elif model_name == 'tf_efficientnet_b3_ns':
self.input_size = (300, 300)
elif model_name == 'tf_efficientnet_b4_ns':
scaleelf.input_size = (380, 380)
elif model_name == 'tf_efficientnet_b6_ns':
self.input_size = (528, 528)
else:
raise Exception('non-valid model name')
# Compose transforms
transform = []
transform += [transforms.Resize(self.input_size)]
self.transform = transforms.Compose(transform)
self.test_dataset = BengaliDataset(self.test_csv,
self.unique_csv,
self.dataset_path,
self.transform,
cache=True)
self.names = self.test_dataset.names
self.test_dataloader = DataLoader(self.test_dataset,
batch_size=self.batch_size,
num_workers=0,
shuffle=False)
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.model_root = MyModel(self.input_size,
self.model_name,
168,
pretrained=True,
dropout=0).to('cuda')
self.model_consonant = MyModel(self.input_size,
self.model_name,
11,
pretrained=True,
dropout=0).to('cuda')
self.model_vowel = MyModel(self.input_size,
self.model_name,
18,
pretrained=True,
dropout=0).to('cuda')
self.model_multihead = MultiHeadModel(self.input_size,
self.model_name,
pretrained=True,
dropout=0).to('cuda')
ckpt = torch.load(self.ckpt)
self.model_root.load_state_dict(ckpt['model_root_state_dict'])
self.model_consonant.load_state_dict(
ckpt['model_consonant_state_dict'])
self.model_vowel.load_state_dict(ckpt['model_vowel_state_dict'])
self.model_multihead.load_state_dict(
ckpt['model_multihead_state_dict'])
def test(self):
pbar = tqdm.tqdm(self.test_dataloader)
pbar.set_description('testing process')
self.model_root.eval()
self.model_consonant.eval()
self.model_vowel.eval()
self.model_multihead.eval()
output_roots = []
output_consonants = []
output_vowels = []
count = 0
with torch.no_grad():
for it, data in enumerate(pbar):
inputs = data[0].to(self.device)
inputs = inputs.repeat(1, 3, 1, 1)
roots = data[1].to(self.device).long()
consonants = data[2].to(self.device).long()
vowels = data[3].to(self.device).long()
uniques = data[4].to(self.device).long()
root_preds, root_preds_2 = self.model_root(inputs, roots)
self.model_root.zero_grad()
consonant_preds, consonant_preds_2 = self.model_consonant(
inputs, consonants)
self.model_consonant.zero_grad()
vowel_preds, vowel_preds_2 = self.model_vowel(inputs, vowels)
self.model_vowel.zero_grad()
root, consonant, vowel, unique, root2, consonant2, vowel2, unique2 = self.model_multihead(
inputs, roots, consonants, vowels, uniques)
self.model_multihead.zero_grad()
unique_prob = F.softmax(unique, dim=1)
for index in range(inputs.shape[0]):
#print('unique:', unique_prob.max(-1).values[index])
if unique_prob.max(-1).values[index] > 0.5:
output_roots.append(root.argmax(-1)[index].item())
output_consonants.append(
consonant.argmax(-1)[index].item())
output_vowels.append(vowel.argmax(-1)[index].item())
else:
output_roots.append(
root_preds.argmax(-1)[index].item())
output_consonants.append(
consonant_preds.argmax(-1)[index].item())
output_vowels.append(
vowel_preds.argmax(-1)[index].item())
row_id, target = [], []
for iid, r, v, c in zip(self.names, output_roots, output_consonants,
output_vowels):
row_id.append(iid + '_grapheme_root')
target.append(int(r))
row_id.append(iid + '_vowel_diacritic')
target.append(int(v))
row_id.append(iid + '_consonant_diacritic')
target.append(int(c))
count += 1
sub_fn = self.output_dir
sub = pd.DataFrame({'row_id': row_id, 'target': target})
sub.to_csv(sub_fn, index=False)
print(f'Done wrote to {sub_fn}')
if args.resume is not None:
name = args.resume
else:
cur_datetime = datetime.now().strftime("%m-%d-%H:%M:%S")
if args.name:
name = "%s-%s" % (cur_datetime, args.name)
else:
name = cur_datetime
model = MyModel(name, lr=args.lr, attn=args.attn, reg=args.reg).to(device)
if args.resume is not None:
ckpt_path = os.path.join('/usr/cs/grad/doc/melgaar/ckpt',
"%s.pt" % args.resume)
state = torch.load(ckpt_path)
model.load_state_dict(state['model_state_dict'])
model.optimizer.load_state_dict(state['optimizer_state_dict'])
model.lr_decay.load_state_dict(state['scheduler'])
model.iter = state['iter']
writer = SummaryWriter(logdir='runs/' + name, purge_step=model.iter)
print("Loaded model")
else:
writer = SummaryWriter(logdir='runs/' + name)
def combine_avg(preds):
return np.expand_dims(np.mean(np.concatenate(preds), 0), 0)
def combine_voting(preds):
preds = np.argmax(np.concatenate(preds), 1)
# if running on GPU and we want to use cuda move model there
use_cuda = torch.cuda.is_available()
if use_cuda:
net = net.cuda()
# create optimizers
optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
# load checkpoint if needed/ wanted
start_n_iter = 0
start_epoch = 0
if opt.resume:
ckpt = load_checkpoint(opt.path_to_checkpoint
) # custom method for loading last checkpoint
net.load_state_dict(ckpt['net'])
start_epoch = ckpt['epoch']
start_n_iter = ckpt['n_iter']
optim.load_state_dict(ckpt['optim'])
print("last checkpoint restored")
# typically we use tensorboardX to keep track of experiments
writer = SummaryWriter()
# now we start the main loop
n_iter = start_n_iter
for epoch in range(start_epoch, opt.epochs):
# set models to train mode
net.train()
# use prefetch_generator and tqdm for iterating through data
model = nn.DataParallel(model)
model.cuda()
optimizer = AdamW(model.parameters(), lr=opt.lr)
log = open(opt.log_path, "a")
max_r = 0
for epoch in range(opt.nepoch):
train_loss = train(model, trainloader, optimizer, opt)
dev_mae, dev_r = test(model, devloader, opt)
print("epoch: %d train_loss: %f dev_mae: %.3f dev_r: %.3f" %
(epoch, train_loss, dev_mae, dev_r))
print("epoch: %d train_loss: %f dev_mae: %.3f dev_r: %.3f" %
(epoch, train_loss, dev_mae, dev_r),
file=log)
if dev_r > max_r:
max_r = dev_r
torch.save(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict()
}, opt.model_path)
checkpoint = torch.load(opt.model_path)
model.load_state_dict(checkpoint['model_state_dict'])
test_mae, test_r = test(model, testloader, opt, filepath=opt.output_path)
print("best epoch: %d test_mae: %.3f test_r: %.3f" %
(checkpoint['epoch'], test_mae, test_r))
acc = correct / tot
print(f'Accuracy: {acc * 100}%')
print(f'Class Accuracy:')
print(class_acc)
print('Confusion Matrix:')
print(confusion)
ALL = args.all
if __name__ == '__main__':
print(f'Begin testing of {MODEL.__class__.__name__}')
if ALL:
i = 0
model_path = os.path.join(MODEL_SAVE_DIR, f'model_epoch_{i}.dict')
while os.path.exists(model_path):
MODEL.load_state_dict(torch.load(model_path))
test(MODEL, TEST_DATALOADER, MODEL_SAVE_DIR)
i += 1
model_path = os.path.join(MODEL_SAVE_DIR, f'model_epoch_{i}.dict')
else:
i = 0
model_path = os.path.join(MODEL_SAVE_DIR, f'model_epoch_{i}.dict')
while os.path.exists(model_path):
i += 1
model_path = os.path.join(MODEL_SAVE_DIR, f'model_epoch_{i}.dict')
model_path = os.path.join(MODEL_SAVE_DIR, f'model_epoch_{i-1}.dict')
test(MODEL, TEST_DATALOADER, MODEL_SAVE_DIR)
test_path = '../test'
total_train_num = 0
total_test_num = 0
for label in os.listdir(train_path):
classes.add(label)
image_num = len(os.listdir(os.path.join(train_path,label)))
total_train_num += image_num
print('train dataset size : {} -> {}'.format(label,image_num))
for label in os.listdir(test_path):
image_num = len(os.listdir(os.path.join(test_path,label)))
total_test_num += image_num
print('test dataset size : {} -> {}'.format(label,image_num))
print('total train dataset : {} \t total test dataset : {}'.format(total_train_num, total_test_num))
return classes
classes = get_classes()
validation_transform = create_validation_transform(True)
test_dataset = MyDataset(transform = validation_transform, mode = 'test', classes = classes)
model = MyModel(num_classes = len(classes))
model.load_state_dict(torch.load(save_path))
model = model.to(device)
test_dataloader = DataLoader(test_dataset,\
batch_size=2,\
shuffle = False
)
test(model, test_dataloader, None, None, device)
pred, trans_feat = model(points)
loss = criterion(pred, target.long())
pred_choice = pred.data.max(1)[1]
correct = pred_choice.eq(target.long().data).cpu().sum()
mean_correct.append(correct.item() / float(len(target)))
mean_loss.append(loss.detach().numpy())
loss.backward()
optimizer.step()
lr_scheduler.step()
epoch_acc = np.mean(mean_correct[-len(dataloader):])
epoch_loss = np.mean(mean_loss[-len(dataloader):])
print(f'Epoch Training Accuracy: {epoch_acc} \t Loss {epoch_loss}')
np.save(os.path.join(save_dir, f'loss_epoch_{epoch}.npy'), (epoch_acc, epoch_loss))
torch.save(model.state_dict(), os.path.join(save_dir,f'model_epoch_{epoch}.dict'))
if __name__=='__main__':
E0 = 0
if os.path.exists(MODEL_SAVE_DIR):
while os.path.exists(os.path.join(MODEL_SAVE_DIR, f'model_epoch_{E0}.dict')):
E0 += 1
print(E0)
if E0 > 0:
MODEL.load_state_dict(torch.load(os.path.join(MODEL_SAVE_DIR, f'model_epoch_{E0-1}.dict')))
print(f'Begin training of {MODEL.__class__.__name__}')
train(MODEL, OPTIMIZER, LR_SCHEDULER, LOSS_CRITERION, TRAIN_DATALOADER, MODEL_SAVE_DIR, E0, EPOCHS)
# loss criterion is cross entropy
len(dataset.nl),
dataset.nl.word_to_idx['<PAD>'],
nn.BatchNorm2d,
num_colors=len(CityFlowNLDataset.colors),
num_types=len(CityFlowNLDataset.vehicle_type) - 2).cuda()
loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=4)
uuids, nls = query(cfg)
saved_dict = torch.load(f'save/{epoch}.pth')
n = {}
for k, v in saved_dict.items():
n[k.replace('module.', '')] = v
model.load_state_dict(n, False)
model.eval()
if os.path.exists('results'):
shutil.rmtree('results')
os.mkdir('results')
# extract img fts first to save time
if not os.path.exists('cache'):
# shutil.rmtree('cache')
os.mkdir('cache')
if not os.path.exists(f'cache/{epoch}'):
os.mkdir(f'cache/{epoch}')
with torch.no_grad():
parser.add_argument("--test_path")
parser.add_argument("--test_batch", default=10, type=int)
parser.add_argument('--max_len', default=512, type=int)
parser.add_argument("--threshold", type=int, default=-1)
parser.add_argument("--pretrained_model_path", default="")
parser.add_argument("--amp", action='store_true')
args = parser.parse_args()
model_dir, file = os.path.split(args.checkpoint_path)
config = pickle.load(open(os.path.join(model_dir, 'args'), 'rb'))
checkpoint = torch.load(os.path.join(model_dir, file),
map_location=torch.device("cpu"))
model_state_dict = checkpoint['model_state_dict']
#replace default vaule if given
config.pretrained_model_path = args.pretrained_model_path if args.pretrained_model_path else config.pretrained_model_path
config.threshold = args.threshold if args.threshold == -1 else config.threshold
mymodel = MyModel(config)
mymodel.load_state_dict(model_state_dict, strict=False)
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
mymodel.to(device)
test_dataloader = load_t1_data(config.dataset_tag, args.test_path,
config.pretrained_model_path,
args.window_size, args.overlap,
args.test_batch, args.max_len)
(p1, r1, f1), (p2, r2, f2) = test_evaluation(mymodel, test_dataloader,
config.threshold, args.amp)
print("Turn 1: precision:{:.4f} recall:{:.4f} f1:{:.4f}".format(
p1, r1, f1))
print("Turn 2: precision:{:.4f} recall:{:.4f} f1:{:.4f}".format(
p2, r2, f2))
def extract_cache_features(cfg, epoch, loader, dataset, test_batch_size, uuids,
nls):
# extract img fts first to save time
if not os.path.exists('cache'):
# shutil.rmtree('cache')
os.mkdir('cache')
if not os.path.exists(f'cache/{epoch}'):
os.mkdir(f'cache/{epoch}')
model = MyModel(cfg,
len(dataset.nl),
dataset.nl.word_to_idx['<PAD>'],
nn.BatchNorm2d,
num_colors=len(CityFlowNLDataset.colors),
num_types=len(CityFlowNLDataset.vehicle_type) -
2).cuda()
saved_dict = torch.load(f'save/{epoch}.pth')
n = {}
for k, v in saved_dict.items():
n[k.replace('module.', '')] = v
model.load_state_dict(n, False)
model.eval()
with torch.no_grad():
for idx, (id, frames, _, _, _, labels) in enumerate(tqdm(loader)):
frames = frames.squeeze(0).cuda()
labels = labels.squeeze(0).cuda()
# b = frames.shape[0]
# cache = []
# version 3
# if b <= test_batch_size:
# cache = model.cnn(frames)
# torch.save(cache, f'cache/{epoch}/{idx}_0.pth')
# else:
# cache = []
for i, (f, l) in enumerate(
zip(frames.split(test_batch_size),
labels.split(test_batch_size))):
cache = model(None, f, l)
img_ft = cache[0]
color = F.softmax(cache[1].mean(dim=0),
dim=0).cpu() #.numpy()
typ = F.softmax(cache[2].mean(dim=0),
dim=0).cpu() #.numpy()
# color = np.argmax(color)
# typ = np.argmax(typ)
cache = [img_ft, color, typ]
torch.save(cache, f'cache/{epoch}/{idx}_{i}.pth')
break
print('saving language features..')
for uuid, query_nl in zip(uuids, nls):
nls_list = []
query_nl, vehicle_type = CityFlowNLDataset.type_replacer(
query_nl)
query_nl, vehicle_color = CityFlowNLDataset.color_replacer(
query_nl)
# max_len = max([len(dataset.nl.do_clean(nl)) for nl in query_nl])
for nl in query_nl:
nl = torch.tensor(
dataset.nl.sentence_to_index(nl,
is_train=False)).cuda()
# nls.append(nl.unsqueeze(0).transpose(1, 0))
nl = nl.unsqueeze(0).transpose(1, 0)
# bs, len, dim
nl = model.rnn(nl)
nls_list.append(nl)
saved_nls = {
'nls': nls_list,
'type': vehicle_type,
'color': vehicle_color
}
torch.save(saved_nls, f'cache/{epoch}/{uuid}.pth')
# model = model.cpu()
del model, saved_dict, n, nls_list, img_ft
torch.cuda.empty_cache()
def train_model_on_dataset(rank, cfg):
dist_rank = rank
# print(dist_rank)
dist.init_process_group(backend="nccl", rank=dist_rank,
world_size=cfg.num_gpu,
init_method="env://")
torch.cuda.set_device(rank)
cudnn.benchmark = True
dataset = CityFlowNLDataset(cfg, build_transforms(cfg))
model = MyModel(cfg, len(dataset.nl), dataset.nl.word_to_idx['<PAD>'], norm_layer=nn.SyncBatchNorm, num_colors=len(CityFlowNLDataset.colors), num_types=len(CityFlowNLDataset.vehicle_type) - 2).cuda()
model = DistributedDataParallel(model, device_ids=[rank],
output_device=rank,
broadcast_buffers=cfg.num_gpu > 1, find_unused_parameters=False)
optimizer = torch.optim.Adam(
params=model.parameters(),
lr=cfg.TRAIN.LR.BASE_LR, weight_decay=0.00003)
lr_scheduler = WarmupMultiStepLR(optimizer,
milestones=cfg.TRAIN.STEPS,
gamma=cfg.TRAIN.LR.WEIGHT_DECAY,
warmup_factor=cfg.TRAIN.WARMUP_FACTOR,
warmup_iters=cfg.TRAIN.WARMUP_EPOCH)
color_loss = LabelSmoothingLoss(len(dataset.colors), 0.1)
vehicle_loss = LabelSmoothingLoss(len(dataset.vehicle_type) - 2, 0.1)
if cfg.resume_epoch > 0:
model.load_state_dict(torch.load(f'save/{cfg.resume_epoch}.pth'))
optimizer.load_state_dict(torch.load(f'save/{cfg.resume_epoch}_optim.pth'))
lr_scheduler.last_epoch = cfg.resume_epoch
lr_scheduler.step()
if rank == 0:
print(f'resume from {cfg.resume_epoch} pth file, starting {cfg.resume_epoch+1} epoch')
cfg.resume_epoch += 1
# loader = DataLoader(dataset, batch_size=cfg.TRAIN.BATCH_SIZE, shuffle=True, num_workers=cfg.TRAIN.NUM_WORKERS)
train_sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=cfg.TRAIN.BATCH_SIZE //cfg.num_gpu,
num_workers=cfg.TRAIN.NUM_WORKERS // cfg.num_gpu,# shuffle=True,
sampler=train_sampler, pin_memory=True)
for epoch in range(cfg.resume_epoch, cfg.TRAIN.EPOCH):
losses = 0.
losses_color = 0.
losses_types = 0.
losses_nl_color = 0.
losses_nl_types = 0.
precs = 0.
train_sampler.set_epoch(epoch)
for idx, (nl, frame, label, act_map, color_label, type_label, nl_color_label, nl_type_label) in enumerate(loader):
# print(nl.shape)
# print(global_img.shape)
# print(local_img.shape)
nl = nl.cuda(non_blocking=True)
label = label.cuda(non_blocking=True)
act_map = act_map.cuda(non_blocking=True)
# global_img, local_img = global_img.cuda(), local_img.cuda()
nl = nl.transpose(1, 0)
frame = frame.cuda(non_blocking=True)
color_label = color_label.cuda(non_blocking=True)
type_label = type_label.cuda(non_blocking=True)
nl_color_label = nl_color_label.cuda(non_blocking=True)
nl_type_label = nl_type_label.cuda(non_blocking=True)
output, color, types, nl_color, nl_types = model(nl, frame, act_map)
# loss = sampling_loss(output, label, ratio=5)
# loss = F.binary_cross_entropy_with_logits(output, label)
total_num_pos = reduce_sum(label.new_tensor([label.sum()])).item()
num_pos_avg_per_gpu = max(total_num_pos / float(cfg.num_gpu), 1.0)
loss = sigmoid_focal_loss(output, label, reduction='sum') / num_pos_avg_per_gpu
loss_color = color_loss(color, color_label) * cfg.TRAIN.ALPHA_COLOR
loss_type = vehicle_loss(types, type_label) * cfg.TRAIN.ALPHA_TYPE
loss_nl_color = color_loss(nl_color, nl_color_label) * cfg.TRAIN.ALPHA_NL_COLOR
loss_nl_type = vehicle_loss(nl_types, nl_type_label) * cfg.TRAIN.ALPHA_NL_TYPE
loss_total = loss + loss_color + loss_type + loss_nl_color + loss_nl_type
optimizer.zero_grad()
loss_total.backward()
# torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
losses += loss.item()
losses_color += loss_color.item()
losses_types += loss_type.item()
losses_nl_color += loss_nl_color.item()
losses_nl_types += loss_nl_type.item()
# precs += recall.item()
if rank == 0 and idx % cfg.TRAIN.PRINT_FREQ == 0:
pred = (output.sigmoid() > 0.5)
# print((pred == label).sum())
pred = (pred == label)
recall = (pred * label).sum() / label.sum()
ca = (color.argmax(dim=1) == color_label)
ca = ca.sum().item() / ca.numel()
ta = (types.argmax(dim=1) == type_label)
ta = ta.sum().item() / ta.numel()
# accu = pred.sum().item() / pred.numel()
lr = optimizer.param_groups[0]['lr']
print(f'epoch: {epoch},',
f'lr: {lr}, step: {idx}/{len(loader)},',
f'loss: {losses / (idx + 1):.4f},',
f'loss color: {losses_color / (idx + 1):.4f},',
f'loss type: {losses_types / (idx + 1):.4f},',
f'loss nl color: {losses_nl_color / (idx + 1):.4f},',
f'loss nl type: {losses_nl_types / (idx + 1):.4f},',
f'recall: {recall.item():.4f}, c_accu: {ca:.4f}, t_accu: {ta:.4f}')
lr_scheduler.step()
if rank == 0:
if not os.path.exists('save'):
os.mkdir('save')
torch.save(model.state_dict(), f'save/{epoch}.pth')
torch.save(optimizer.state_dict(), f'save/{epoch}_optim.pth')
