class DummyProbe(Behavior):
def __init__(self, logger, params):
self.lick_timer = Timer()
self.lick_timer.start()
self.ready_timer = Timer()
self.ready_timer.start()
self.ready = False
self.probe = 0
super(DummyProbe, self).__init__(logger, params)
def is_ready(self, init_duration):
self.__get_events()
elapsed_time = self.ready_timer.elapsed_time()
return self.ready and elapsed_time > init_duration
def inactivity_time(self): # in minutes
return self.lick_timer.elapsed_time() / 1000 / 60
def is_licking(self, since=0):
probe = self.__get_events()
# reset lick timer if licking is detected &
if probe > 0:
self.resp_timer.start()
self.licked_probe = probe
return probe
def __get_events(self):
probe = 0
events = pygame.event.get()
for event in events:
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_LEFT:
self.logger.log_lick(1)
print('Probe 1 activated!')
probe = 1
self.lick_timer.start()
elif event.key == pygame.K_RIGHT:
self.logger.log_lick(2)
print('Probe 2 activated!')
probe = 2
elif event.key == pygame.K_SPACE and self.ready:
self.ready = False
print('off position')
elif event.key == pygame.K_SPACE and not self.ready:
self.lick_timer.start()
self.ready = True
print('in position')
return probe
class Agent(Panda):
def __init__(self, env, cond):
self.cond = cond
self.env = env
self.timer = Timer()
self.duration = cond['dur']
hfov = self.env.camLens.get_fov(
) * np.pi / 180 # half field of view in radians
# define object time parameters
self.rot_fun = self.time_fun(cond['rot'])
self.tilt_fun = self.time_fun(cond['tilt'])
self.yaw_fun = self.time_fun(cond['yaw'])
z_loc = 2
self.x_fun = self.time_fun(cond['pos_x'],
lambda x, t: np.arctan(x * hfov[0]) * z_loc)
self.y_fun = self.time_fun(cond['pos_y'],
lambda x, t: np.arctan(x * hfov[0]) * z_loc)
self.scale_fun = self.time_fun(cond['mag'], lambda x, t: .15 * x)
# add task object
self.name = "Obj%s-Task" % cond['id']
def run(self):
self.model = self.env.loader.loadModel(
self.env.object_files[self.cond['id']])
self.model.reparentTo(self.env.render)
self.task = self.env.taskMgr.doMethodLater(self.cond['delay'] / 1000,
self.objTask, self.name)
def objTask(self, task):
t = self.timer.elapsed_time() / 1000
if t > self.duration / 1000:
self.remove(task)
return
self.model.setHpr(self.rot_fun(t), self.tilt_fun(t), self.yaw_fun(t))
self.model.setPos(self.x_fun(t), 2, self.y_fun(t))
self.model.setScale(self.scale_fun(t))
return Task.cont
def remove(self, task):
task.remove()
self.model.removeNode()
def time_fun(self, param, fun=lambda x, t: x):
param = (iterable(param))
idx = np.linspace(0, self.duration / 1000, param.size)
return lambda t: np.interp(t, idx, fun(param, t))
class DummyPorts(MultiPort):
def setup(self, exp):
import pygame
self.lick_timer = Timer()
self.lick_timer.start()
self.ready_timer = Timer()
self.ready_timer.start()
self.ready = False
self.interface = 0
pygame.init()
#self.screen = pygame.display.set_mode((800, 480))
self.params = exp.params
self.resp_timer = Timer()
self.resp_timer.start()
self.logger = exp.logger
self.exp = exp
self.rew_port = 0
self.choices = np.array(np.empty(0))
self.choice_history = list() # History term for bias calculation
self.reward_history = list(
) # History term for performance calculation
self.licked_port = 0
self.reward_amount = dict()
self.curr_cond = []
def is_ready(self, duration, since=0):
if duration == 0: return True
self.__get_events()
elapsed_time = self.ready_timer.elapsed_time()
return self.ready and elapsed_time >= duration
def is_licking(self, since=0):
port = self.__get_events()
if port > 0: self.resp_timer.start()
self.licked_port = port
return port
def get_response(self, since=0):
port = self.is_licking(since)
return port > 0
def is_correct(self):
return np.any(
np.equal(self.licked_port, self.curr_cond['response_port']))
def reward(self):
self.update_history(self.licked_port, self.reward_amount)
self.log_reward(self.reward_amount)
print('Giving Water at port:%1d' % self.licked_port)
return True
def prepare(self, condition):
self.curr_cond = condition
self.reward_amount = condition['reward_amount']
self.logger.log('BehCondition.Trial',
dict(beh_hash=self.curr_cond['beh_hash']),
schema='behavior')
def punish(self):
print('punishing')
port = self.licked_port if self.licked_port > 0 else np.nan
self.update_history(port)
def exit(self):
pass
def __get_events(self):
port = 0
events = pygame.event.get()
for event in events:
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_LEFT:
self.logger.log('Activity.Lick',
dict(port=1),
schema='behavior')
print('Probe 1 activated!')
port = 1
self.lick_timer.start()
elif event.key == pygame.K_RIGHT:
self.logger.log('Activity.Lick',
dict(port=2),
schema='behavior')
print('Probe 2 activated!')
port = 2
elif event.key == pygame.K_SPACE and not self.ready:
self.lick_timer.start()
self.ready = True
self.log_activity('Proximity',
dict(port=3, in_position=self.ready))
print('in position')
elif event.type == pygame.KEYUP:
if event.key == pygame.K_SPACE and self.ready:
self.ready = False
self.log_activity('Proximity',
dict(port=3, in_position=self.ready))
print('off position')
print(pygame.mouse.get_pos())
elif event.type == pygame.MOUSEBUTTONDOWN:
print(pygame.mouse.get_pos())
return port
class Logger:
def __init__(self, protocol=False):
self.setup, self.is_pi = socket.gethostname(), os.uname(
)[4][:3] == 'arm'
self.curr_state, self.lock, self.queue, self.curr_trial, self.total_reward = '', False, PriorityQueue(
), 0, 0
self.session_key = dict()
self.ping_timer, self.logger_timer = Timer(), Timer()
self.setup_status = 'running' if protocol else 'ready'
self.log_setup(protocol)
fileobject = open(
os.path.dirname(os.path.abspath(__file__)) + '/dj_local_conf.json')
connect_info = json.loads(fileobject.read())
background_conn = dj.Connection(connect_info['database.host'],
connect_info['database.user'],
connect_info['database.password'])
self.schemata = dict()
self.schemata['lab'] = dj.create_virtual_module(
'beh.py', 'lab_behavior', connection=background_conn)
self.schemata['mice'] = dj.create_virtual_module(
'mice.py', 'lab_mice', connection=background_conn)
self.thread_end, self.thread_lock = threading.Event(), threading.Lock()
self.inserter_thread = threading.Thread(target=self.inserter)
self.getter_thread = threading.Thread(target=self.getter)
self.inserter_thread.start()
self.getter_thread.start()
self.logger_timer.start() # start session time
def put(self, **kwargs):
self.queue.put(PrioritizedItem(**kwargs))
def inserter(self):
while not self.thread_end.is_set():
if self.queue.empty():
time.sleep(.5)
continue
item = self.queue.get()
ignore, skip = (False, False) if item.replace else (True, True)
table = self.rgetattr(self.schemata[item.schema], item.table)
self.thread_lock.acquire()
table.insert1(item.tuple,
ignore_extra_fields=ignore,
skip_duplicates=skip,
replace=item.replace)
self.thread_lock.release()
def getter(self):
while not self.thread_end.is_set():
self.thread_lock.acquire()
self.setup_info = (self.schemata['lab'].SetupControl()
& dict(setup=self.setup)).fetch1()
self.thread_lock.release()
self.setup_status = self.setup_info['status']
time.sleep(1) # update once a second
def log(self, table, data=dict()):
tmst = self.logger_timer.elapsed_time()
self.put(table=table,
tuple={
**self.session_key, 'trial_idx': self.curr_trial,
'time': tmst,
**data
})
return tmst
def log_setup(self, task_idx=False):
rel = SetupControl() & dict(setup=self.setup)
key = rel.fetch1() if numpy.size(rel.fetch()) else dict(
setup=self.setup)
if task_idx: key['task_idx'] = task_idx
key = {**key, 'ip': self.get_ip(), 'status': self.setup_status}
SetupControl.insert1(key, replace=True)
def log_session(self, params, exp_type=''):
self.curr_trial, self.total_reward, self.session_key = 0, 0, {
'animal_id': self.get_setup_info('animal_id')
}
last_sessions = (Session() & self.session_key).fetch('session')
self.session_key['session'] = 1 if numpy.size(
last_sessions) == 0 else numpy.max(last_sessions) + 1
self.put(table='Session',
tuple={
**self.session_key, 'session_params': params,
'setup': self.setup,
'protocol': self.get_protocol(),
'experiment_type': exp_type
},
priority=1)
key = {
'session': self.session_key['session'],
'trials': 0,
'total_liquid': 0,
'difficulty': 1
}
if 'start_time' in params:
tdelta = lambda t: datetime.strptime(
t, "%H:%M:%S") - datetime.strptime("00:00:00", "%H:%M:%S")
key = {
**key, 'start_time': tdelta(params['start_time']),
'stop_time': tdelta(params['stop_time'])
}
self.update_setup_info(key)
self.logger_timer.start() # start session time
def log_conditions(self, conditions, condition_tables=[]):
for cond in conditions:
cond_hash = make_hash(cond)
self.put(table='Condition',
tuple=dict(cond_hash=cond_hash, cond_tuple=cond.copy()),
priority=5)
cond.update({'cond_hash': cond_hash})
for condtable in condition_tables:
if condtable == 'RewardCond' and isinstance(
cond['probe'], tuple):
for idx, probe in enumerate(cond['probe']):
self.put(table=condtable,
tuple={
'cond_hash': cond['cond_hash'],
'probe': probe,
'reward_amount': cond['reward_amount']
})
else:
self.put(table=condtable, tuple=dict(cond.items()))
if condtable == 'OdorCond':
for idx, port in enumerate(cond['delivery_port']):
self.put(table=condtable + '.Port',
tuple={
'cond_hash': cond['cond_hash'],
'dutycycle': cond['dutycycle'][idx],
'odor_id': cond['odor_id'][idx],
'delivery_port': port
})
return conditions
def init_trial(self, cond_hash):
self.curr_trial += 1
if self.lock: self.thread_lock.acquire()
self.curr_cond, self.trial_start = cond_hash, self.logger_timer.elapsed_time(
)
return self.trial_start # return trial start time
def log_trial(self, last_flip_count=0):
if self.lock: self.thread_lock.release()
timestamp = self.logger_timer.elapsed_time()
self.put(table='Trial',
tuple=dict(self.session_key,
trial_idx=self.curr_trial,
cond_hash=self.curr_cond,
start_time=self.trial_start,
end_time=timestamp,
last_flip_count=last_flip_count))
def log_pulse_weight(self, pulse_dur, probe, pulse_num, weight=0):
key = dict(setup=self.setup,
probe=probe,
date=systime.strftime("%Y-%m-%d"))
self.put(table='LiquidCalibration', tuple=key, priority=5)
self.put(table='LiquidCalibration.PulseWeight',
tuple=dict(key,
pulse_dur=pulse_dur,
pulse_num=pulse_num,
weight=weight),
replace=True)
def update_setup_info(self, info):
self.setup_info = {
**(SetupControl() & dict(setup=self.setup)).fetch1(),
**info
}
self.put(table='SetupControl',
tuple=self.setup_info,
replace=True,
priority=1)
self.setup_status = self.setup_info['status']
if 'status' in info:
while self.get_setup_info('status') != info['status']:
time.sleep(.5)
def get_setup_info(self, field):
return (SetupControl() & dict(setup=self.setup)).fetch1(field)
def get_protocol(self, task_idx=None):
if not task_idx: task_idx = self.get_setup_info('task_idx')
if len(Task() & dict(task_idx=task_idx)) > 0:
protocol = (Task() & dict(task_idx=task_idx)).fetch1('protocol')
path, filename = os.path.split(protocol)
if not path:
protocol = str(pathlib.Path(
__file__).parent.absolute()) + '/conf/' + filename
return protocol
else:
return False
def ping(self, period=5000):
if self.ping_timer.elapsed_time(
) >= period: # occasionally update control table
self.ping_timer.start()
self.update_setup_info({
'last_ping':
str(datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
'queue_size':
self.queue.qsize(),
'trials':
self.curr_trial,
'total_liquid':
self.total_reward,
'state':
self.curr_state
})
def cleanup(self):
while not self.queue.empty():
print('Waiting for empty queue... qsize: %d' % self.queue.qsize())
time.sleep(2)
self.thread_end.set()
@staticmethod
def get_ip():
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
s.connect(("8.8.8.8", 80))
IP = s.getsockname()[0]
except Exception:
IP = '127.0.0.1'
finally:
s.close()
return IP
@staticmethod
def rgetattr(obj, attr, *args):
def _getattr(obj, attr):
return getattr(obj, attr, *args)
return functools.reduce(_getattr, [obj] + attr.split('.'))
class Logger:
""" This class handles the database logging"""
def __init__(self):
self.last_trial = 0
self.queue = Queue()
self.timer = Timer()
self.trial_start = 0
self.curr_cond = []
self.task_idx = []
self.session_key = dict()
self.setup = socket.gethostname()
self.lock = True
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.connect(("8.8.8.8", 80))
self.ip = s.getsockname()[0]
path = os.path.dirname(os.path.abspath(__file__))
fileobject = open(path + '/dj_local_conf.json')
connect_info = json.loads(fileobject.read())
conn2 = dj.Connection(connect_info['database.host'],
connect_info['database.user'],
connect_info['database.password'])
self.insert_schema = dj.create_virtual_module('beh.py',
'lab_behavior',
connection=conn2)
self.thread_end = threading.Event()
self.thread_lock = Lock()
self.thread_runner = threading.Thread(
target=self.inserter) # max insertion rate of 10 events/sec
self.thread_runner.start()
def cleanup(self):
self.thread_end.set()
def inserter(self):
while not self.thread_end.is_set():
if not self.queue.empty():
self.thread_lock.acquire()
item = self.queue.get()
if 'update' in item:
eval(
'(self.insert_schema.' + item['table'] +
'() & item["tuple"])._update(item["field"],item["value"])'
)
else:
eval(
'self.insert_schema.' + item['table'] +
'.insert1(item["tuple"], ignore_extra_fields=True, skip_duplicates=True)'
)
self.thread_lock.release()
else:
time.sleep(.5)
def log_setup(self):
key = dict(setup=self.setup)
# update values in case they exist
if numpy.size((SetupControl() & dict(setup=self.setup)).fetch()):
key = (SetupControl() & dict(setup=self.setup)).fetch1()
(SetupControl() & dict(setup=self.setup)).delete_quick()
# insert new setup
key['ip'] = self.ip
key['status'] = 'ready'
SetupControl().insert1(key)
def log_session(self, session_params, exp_type=''):
animal_id, task_idx = (SetupControl() & dict(setup=self.setup)).fetch1(
'animal_id', 'task_idx')
self.task_idx = task_idx
self.last_trial = 0
# create session key
self.session_key = dict()
self.session_key['animal_id'] = animal_id
last_sessions = (Session() & self.session_key).fetch('session')
if numpy.size(last_sessions) == 0:
last_session = 0
else:
last_session = numpy.max(last_sessions)
self.session_key['session'] = last_session + 1
# get task parameters for session table
key = dict(self.session_key.items())
key['session_params'] = session_params
key['setup'] = self.setup
key['protocol'] = self.get_protocol()
key['experiment_type'] = exp_type
self.queue.put(dict(table='Session', tuple=key))
# start session time
self.timer.start()
(SetupControl() & dict(setup=self.setup))._update(
'current_session', self.session_key['session'])
(SetupControl() & dict(setup=self.setup))._update('last_trial', 0)
(SetupControl() & dict(setup=self.setup))._update('total_liquid', 0)
def log_conditions(
self,
conditions,
condition_tables=['OdorCond', 'MovieCond', 'RewardCond']):
# iterate through all conditions and insert
for cond in conditions:
cond_hash = make_hash(cond)
self.queue.put(
dict(table='Condition',
tuple=dict(cond_hash=cond_hash, cond_tuple=cond.copy())))
cond.update({'cond_hash': cond_hash})
for condtable in condition_tables:
if condtable == 'RewardCond' and isinstance(
cond['probe'], tuple):
for idx, probe in enumerate(cond['probe']):
key = {
'cond_hash': cond['cond_hash'],
'probe': probe,
'reward_amount': cond['reward_amount']
}
self.queue.put(dict(table=condtable, tuple=key))
else:
self.queue.put(
dict(table=condtable, tuple=dict(cond.items())))
if condtable == 'OdorCond':
for idx, port in enumerate(cond['delivery_port']):
key = {
'cond_hash': cond['cond_hash'],
'dutycycle': cond['dutycycle'][idx],
'odor_id': cond['odor_id'][idx],
'delivery_port': port
}
self.queue.put(
dict(table=condtable + '.Port', tuple=key))
return conditions
def init_trial(self, cond_hash):
self.curr_cond = cond_hash
if self.lock:
self.thread_lock.acquire()
# return trial start time
self.trial_start = self.timer.elapsed_time()
return self.trial_start
def log_trial(self, last_flip_count=0):
if self.lock:
self.thread_lock.release()
timestamp = self.timer.elapsed_time()
trial_key = dict(self.session_key,
trial_idx=self.last_trial + 1,
cond_hash=self.curr_cond,
start_time=self.trial_start,
end_time=timestamp,
last_flip_count=last_flip_count)
self.queue.put(dict(table='Trial', tuple=trial_key))
self.last_trial += 1
# insert ping
self.queue.put(
dict(table='SetupControl',
tuple=dict(setup=self.setup),
field='last_trial',
value=self.last_trial,
update=True))
def log_liquid(self, probe, reward_amount):
timestamp = self.timer.elapsed_time()
self.queue.put(
dict(table='LiquidDelivery',
tuple=dict(self.session_key,
time=timestamp,
probe=probe,
reward_amount=reward_amount)))
def log_stim(self):
timestamp = self.timer.elapsed_time()
self.queue.put(
dict(table='StimOnset',
tuple=dict(self.session_key, time=timestamp)))
def log_lick(self, probe):
timestamp = self.timer.elapsed_time()
self.queue.put(
dict(table='Lick',
tuple=dict(self.session_key, time=timestamp, probe=probe)))
return timestamp
def log_pulse_weight(self, pulse_dur, probe, pulse_num, weight=0):
cal_key = dict(setup=self.setup,
probe=probe,
date=systime.strftime("%Y-%m-%d"))
LiquidCalibration().insert1(cal_key, skip_duplicates=True)
(LiquidCalibration.PulseWeight()
& dict(cal_key, pulse_dur=pulse_dur)).delete_quick()
LiquidCalibration.PulseWeight().insert1(
dict(cal_key,
pulse_dur=pulse_dur,
pulse_num=pulse_num,
weight=weight))
def log_animal_weight(self, weight):
key = dict(animal_id=self.get_setup_info('animal_id'), weight=weight)
Mice.MouseWeight().insert1(key)
def log_position(self, in_position, state):
timestamp = self.timer.elapsed_time()
self.queue.put(
dict(table='CenterPort',
tuple=dict(self.session_key,
time=timestamp,
in_position=in_position,
state=state)))
def update_setup_status(self, status):
key = (SetupControl() & dict(setup=self.setup)).fetch1()
in_status = key['status'] == status
if not in_status:
(SetupControl() & dict(setup=self.setup))._update('status', status)
return in_status
def update_setup_notes(self, note):
self.queue.put(
dict(table='SetupControl',
tuple=dict(setup=self.setup),
field='notes',
value=note,
update=True))
def update_state(self, state):
self.queue.put(
dict(table='SetupControl',
tuple=dict(setup=self.setup),
field='state',
value=state,
update=True))
def update_animal_id(self, animal_id):
(SetupControl() & dict(setup=self.setup))._update(
'animal_id', animal_id)
def update_task_idx(self, task_idx):
(SetupControl() & dict(setup=self.setup))._update('task_idx', task_idx)
def update_total_liquid(self, total_rew):
self.queue.put(
dict(table='SetupControl',
tuple=dict(setup=self.setup),
field='total_liquid',
value=total_rew,
update=True))
def update_difficulty(self, difficulty):
self.queue.put(
dict(table='SetupControl',
tuple=dict(setup=self.setup),
field='difficulty',
value=difficulty,
update=True))
def update_setup_info(self, field, value):
self.queue.put(
dict(table='SetupControl',
tuple=dict(setup=self.setup),
field=field,
value=value,
update=True))
def get_setup_info(self, field):
info = (SetupControl() & dict(setup=self.setup)).fetch1(field)
return info
def get_session_key(self):
return self.session_key
def get_clip_info(self, key):
clip_info = (MovieTables.Movie() * MovieTables.Movie.Clip() & key
& self.session_key).fetch1()
return clip_info
def get_protocol(self):
task_idx = (SetupControl() & dict(setup=self.setup)).fetch1('task_idx')
protocol = (Task() & dict(task_idx=task_idx)).fetch1('protocol')
path, filename = os.path.split(protocol)
if not path:
path = pathlib.Path(__file__).parent.absolute()
protocol = str(path) + '/conf/' + filename
return protocol
def ping(self):
lp = str(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
self.queue.put(
dict(table='SetupControl',
tuple=dict(setup=self.setup),
field='last_ping',
value=lp,
update=True))
self.queue.put(
dict(table='SetupControl',
tuple=dict(setup=self.setup),
field='queue_size',
value=self.queue.qsize(),
update=True))
class DummyProbe(Behavior):
def __init__(self, logger, params):
import pygame
self.lick_timer = Timer()
self.lick_timer.start()
self.ready_timer = Timer()
self.ready_timer.start()
self.ready = False
self.interface = 0
pygame.init()
self.screen = pygame.display.set_mode((800, 480))
super(DummyProbe, self).__init__(logger, params)
def get_cond_tables(self):
return ['RewardCond']
def is_ready(self, duration, since=0):
if duration == 0: return True
self.__get_events()
elapsed_time = self.ready_timer.elapsed_time()
return self.ready and elapsed_time >= duration
def is_licking(self,since=0):
probe = self.__get_events()
if probe > 0: self.resp_timer.start()
self.licked_probe = probe
return probe
def get_response(self, since=0):
probe = self.is_licking(since)
return probe > 0
def is_correct(self):
return np.any(np.equal(self.licked_probe, self.curr_cond['probe']))
def prepare(self, condition):
self.curr_cond = condition
self.reward_amount = condition['reward_amount']
def reward(self):
self.update_history(self.licked_probe, self.reward_amount)
self.logger.log('LiquidDelivery', dict(probe=self.licked_probe,
reward_amount=self.reward_amount))
print('Giving Water at probe:%1d' % self.licked_probe)
return True
def punish(self):
print('punishing')
probe = self.licked_probe if self.licked_probe > 0 else np.nan
self.update_history(probe)
def __get_events(self):
probe = 0
events = pygame.event.get()
for event in events:
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_LEFT:
self.logger.log('Lick', dict(probe=1))
print('Probe 1 activated!')
probe = 1
self.lick_timer.start()
elif event.key == pygame.K_RIGHT:
self.logger.log('Lick', dict(probe=2))
print('Probe 2 activated!')
probe = 2
elif event.key == pygame.K_SPACE and not self.ready:
self.lick_timer.start()
self.ready = True
print('in position')
elif event.type == pygame.KEYUP:
if event.key == pygame.K_SPACE and self.ready:
self.ready = False
print('off position')
print(pygame.mouse.get_pos())
elif event.type == pygame.MOUSEBUTTONDOWN:
print(pygame.mouse.get_pos())
return probe
def main(args):
device = torch.device(args.device)
logger = SummaryWriter(
os.path.join(
args.logs_dir,
"tensorboard_log{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.now())))
if args.dataset not in DATASETS:
raise Exception(
f'`--dataset` is invalid. it should be one of {list(DATASETS.keys())}'
)
train_data = DATASETS[args.dataset](args.train_root,
transforms=Compose([
RandomCrop(224),
RandomHorizontalFlip(),
RandomVerticalFlip(),
ToTensor(),
Normalize([0.5, 0.5, 0.5],
[0.5, 0.5, 0.5])
]),
**args.__dict__)
train_loader = DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=args.num_workers)
do_val = False
if args.val_root:
val_data = DATASETS[args.dataset](args.val_root,
transforms=ToTensor(),
**args.__dict__)
val_loader = DataLoader(val_data,
batch_size=args.batch_size,
num_workers=args.num_workers)
do_val = True
net = MODELS[args.model_name](pretrained=args.resume is None,
**args.__dict__).to(device)
optimizer = torch.optim.SGD(params=net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=1e-4)
start_epoch, total_epoch, global_step = 0, args.epochs, 0
if args.resume is not None:
checkpoint = load_checkpoint(args.resume)
net.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch']
global_step = checkpoint['global_step']
print(f"=> Start epoch {start_epoch} ")
for epoch in range(start_epoch, total_epoch):
net.train()
timer, counter = Timer(), Counter()
timer.start()
for step, (img, label) in enumerate(train_loader):
img, label = img.to(device), label.to(device)
reader_time = timer.elapsed_time()
loss, miou = net(img, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss = float(loss)
batch_time = timer.elapsed_time()
counter.append(loss=loss,
miou=miou,
reader_time=reader_time,
batch_time=batch_time)
eta = calculate_eta(len(train_loader) - step, counter.batch_time)
print(
f"[epoch={epoch + 1}/{total_epoch}] "
f"step={step + 1}/{len(train_loader)} "
f"loss={loss:.4f}/{counter.loss:.4f} "
f"miou={miou:.4f}/{counter.miou:.4f} "
f"batch_time={counter.batch_time:.4f} "
f"reader_time={counter.reader_time:.4f} "
f"| ETA {eta}",
end="\r",
flush=True)
logger.add_scalar("loss", float(loss), global_step=global_step)
logger.add_scalar("miou", float(miou), global_step=global_step)
global_step += 1
logger.flush()
timer.restart()
print()
save_checkpoint(
{
'net': net.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': epoch + 1,
'global_step': global_step,
},
epoch + 1,
False,
save_dir=args.save_dir)
pass
pass
