def init_training_configs(batch_size, model, alpha):
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())
# exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
trainset = utils.SoccerBallDataset("data/train/data.csv", "data/train", downsample=4, alpha=alpha)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)
print("# examples: ", len(trainset))
return criterion, optimizer, trainloader, trainset
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--loadSweaty',
required=True,
help='path to pretrained Sweaty model')
parser.add_argument('--loadGru',
default='',
help='path to pretrained Gru cell')
parser.add_argument('--downsample', type=int, default=4, help='downsample')
parser.add_argument('--alpha',
type=int,
default=1000,
help='multiplication factor for the teacher signals')
parser.add_argument('--seq_len',
type=int,
default=10,
help='length of the sequence')
parser.add_argument('--gruType',
type=int,
default=1,
help='length of the sequence')
parser.add_argument('--cpu',
type=int,
default=0,
help='run inference time in cpu')
opt = parser.parse_args()
gruType = opt.gruType
loadGru = opt.loadGru
if opt.cpu != 0:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
else:
device = torch.device("cpu")
print("Running inference time on ", device)
trainset = utils.SoccerBallDataset("data/train/data.csv",
"data/train",
downsample=4,
alpha=opt.alpha)
if loadGru == '':
sweaty = SweatyNet1()
sweaty.to(device)
sweaty.load_state_dict(torch.load(opt.loadSweaty))
inference_sweaty(sweaty, trainset, device)
elif loadGru != '' and gruType == 0:
print("Loading sweaty and gru for type 1")
sweaty = SweatyNet1()
sweaty.to(device)
sweaty.load_state_dict(torch.load(opt.loadSweaty))
conv_gru = gru.ConvGruCell(1, 1, device=device)
conv_gru.to(device)
conv_gru.load_state_dict(torch.load(loadGru))
inference_gru(sweaty, conv_gru, trainset, device)
else:
print("Loading sweaty and gru for type 2")
sweaty = s2.SweatyNet1()
sweaty.to(device)
sweaty.load_state_dict(torch.load(opt.loadSweaty))
conv_gru = gru.ConvGruCell(89, 1, device=device)
conv_gru.to(device)
conv_gru.load_state_dict(torch.load(loadGru))
inference_type2_gru(sweaty, conv_gru, trainset, device)
print(device)
opt = parser.parse_args()
downsample = opt.downsample
batch_size = opt.batch_size
testset = opt.testSet
trainset = opt.trainSet
pretrained_model = opt.load
model = SweatyNet1()
model.load_state_dict(torch.load(pretrained_model))
model.to(device)
testset = utils.SoccerBallDataset(testset + "data.csv",
testset,
downsample=downsample,
alpha=opt.alpha)
trainset = utils.SoccerBallDataset(trainset + "data.csv",
trainset,
downsample=downsample,
alpha=opt.alpha)
model.eval()
threshold = utils.get_abs_threshold(trainset)
metrics = utils.evaluate_sweaty_model(model,
device,
testset,
threshold,
verbose=True,
debug=False)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
modelPath = ''
model = SweatyNet1()
model.to(device)
if modelPath != '':
print("Loading Sweaty")
model.load_state_dict(torch.load(modelPath, map_location='cpu'))
print(model)
conv_gru = ConvGruCell(1, 1, device)
trainset = utils.SoccerBallDataset("data/train/data.csv",
"data/train",
downsample=4,
delimiter=',')
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=2,
shuffle=True,
num_workers=2)
idx = 51
image = trainset[idx]['image']
signal = np.array(trainset[idx]['signal'].squeeze())
output = model(image.unsqueeze(0).float().to(device))
print(output.size())
output = conv_gru(output)
output_signal = np.array(output.cpu().squeeze().detach())