args.dtype = torch.cuda.FloatTensor if args.cuda else torch.FloatTensor
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed) # Safe even if no GPU
from model import RN, CNN_MLP
from model_vigil import RFES, RFESH
if args.model == 'RFES':
model = RFES(args)
elif args.model == 'RFESH':
model = RFESH(args)
elif args.model == 'CNN_MLP':
model = CNN_MLP(args)
else:
model = RN(args)
print(args)
# For loading the data (possibly a symlink to relational-networks/data)
data_dirs = './data'
bs = args.batch_size
input_img = torch.FloatTensor(bs, 3, 75, 75)
input_qst = torch.FloatTensor(bs, 11)
label = torch.LongTensor(bs)
if args.cuda:
model.cuda()
parser.add_argument(
'--dropout_prob',
type=float,
default=0.0,
help='dropout probability just before the final FCN (default: 0.0)')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
kwargs = dict(args._get_kwargs())
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
if args.model == 'CNN_MLP':
model = CNN_MLP(kwargs)
elif args.model == 'RN':
model = RN(kwargs)
elif args.model == 'RN2':
model = RN2(kwargs)
else:
model = MHDPA_RN(kwargs)
#model_dirs = './model_SimpleMHOutput'
#model_dirs = './model_2LMHOutput'
model_dirs = './model_3LMHOutput'
#model_dirs = './model_4LMHOutput'
bs = args.batch_size
input_img = torch.FloatTensor(bs, 3, 75, 75)
input_qst = torch.FloatTensor(bs, 11)
label = torch.LongTensor(bs)