def get_output_path(args):
if args.model.suffix == '.pt':
model_path = args.model
else:
serializer = Serializer(args.model)
model_path = serializer._id2path(args.step)
return args.output / (model_path.stem + '.pkl')
def get_samples_passed(args):
serializer = Serializer(args.model)
model_path = serializer._id2path(args.step)
data = torch.load(model_path, map_location='cpu')
return data.get('samples_passed', data['global_step'] * args.bs)